"""
=============================================================================
第 11 讲 · 工具模块：参数高效微调（PEFT）的几种做法，从零实现
=============================================================================

对应讲义：lectures/lecture11_efficient_adaptation.md 第 4~7 节
被本目录的 02、03 号文件使用。

-----------------------------------------------------------------------------
【这个模块要解决什么问题？】
-----------------------------------------------------------------------------
全量微调（full fine-tuning）要更新模型的**每一个**参数：要存每个参数的梯度和优化器状态，
每个任务还要单独存一份完整的模型。参数高效微调（parameter-efficient fine-tuning，PEFT）的思路是：
    **把预训练的参数冻结（一个都不动），只训练很少量的新参数（或一小部分旧参数）。**
这里实现讲义里讲到的四种做法，全部作用在第 7 讲的 TransformerLM 上：

    LoRA            LoRALinear / add_lora() / merge_lora()
                    给线性层加一条"低秩旁路"：h = W x + (α/r)·B A x，只训练 A、B
    适配器 Adapter   Adapter / add_adapters()
                    在每个子层的输出后面插一个小的"瓶颈"网络：h + W_up·GELU(W_down·h)
    提示微调         SoftPromptModel
                    在输入序列前面拼几个**可训练的向量**（"软提示"），模型本身一点不动
    BitFit          bitfit()
                    只训练所有的偏置项（bias）

以及几个通用的小工具：freeze_all()、count_trainable()、trainable_state_dict()。

【用法】
    import peft
    model = ...                                  # 一个 TransformerLM（比如第 9 讲预训练好的）
    peft.freeze_all(model)                       # 先把所有参数冻结
    peft.add_lora(model, r=4, alpha=8)           # 再加上 LoRA（新加的 A、B 是可训练的）
    print(peft.count_trainable(model))           # 可训练的参数有多少
    ...训练...
    peft.merge_lora(model)                       # 训练完把 BA 合并进 W，推理时没有额外开销

直接运行本文件会做一组自我检查（约 1 秒）。
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


# =============================================================================
# 第 1 部分：通用小工具
# =============================================================================
def freeze_all(model):
    """
    把模型的所有参数冻结（requires_grad = False），之后优化器不会更新它们，也不会为它们计算梯度。

    【参数】model：任意 nn.Module
    【返回】model 本身（方便链式调用）

    【Python 小知识】
        每个参数（nn.Parameter）都有一个 requires_grad 属性。设为 False 后，
        反向传播时 PyTorch 不会计算它的梯度（省计算、省内存），它的 .grad 一直是 None。
        冻结的参数**仍然参与前向计算**，只是不被更新。
    """
    for p in model.parameters():
        p.requires_grad = False
    return model


def count_trainable(model):
    """
    返回 (可训练的参数个数, 全部参数个数)。

    【边界】共用的参数（例如第 7 讲里输出层和词嵌入共用同一个矩阵）只算一次：
        model.parameters() 会自动去重。
    """
    params = list(model.parameters())
    return sum(p.numel() for p in params if p.requires_grad), sum(p.numel() for p in params)


def trainable_state_dict(model):
    """
    只取出**可训练**参数的 state_dict —— 这就是一个任务需要单独保存的全部内容。

    【返回】dict：参数名 → 张量（复制了一份，之后修改模型不影响它）

    【NLP 联系】这就是 PEFT 最实际的好处之一：一个基座模型 + 每个任务几 MB 的"补丁"。
        LoRA 论文里，GPT-3 175B 的一个完整检查点约 350 GB，而一个 LoRA 补丁（r=4，只加在 Wq、Wv 上）约 35 MB。
    """
    trainable = {id(p) for p in model.parameters() if p.requires_grad}
    return {name: p.detach().clone() for name, p in model.named_parameters() if id(p) in trainable}


def bitfit(model):
    """
    BitFit（Ben Zaken 等 2021）：冻结所有参数，然后只放开名字以 ".bias" 结尾的参数（所有偏置项）。

    【返回】model

    【思路】偏置只是给每个神经元的输出加一个常数，参数极少（我们的迷你 GPT 里不到 2%，大模型里约 0.1%），
        但它能整体地"挪动"每一层的激活值。BitFit 的发现是：在不少分类任务上，只调偏置就能接近全量微调。
    【边界】LayerNorm 也有 bias（和 weight），这里一并放开它的 bias；没有偏置的层（如我们的输出层）不受影响。
    """
    freeze_all(model)
    for name, p in model.named_parameters():
        if name.endswith(".bias"):
            p.requires_grad = True
    return model


def _replace_module(model, full_name, new_module):
    """
    把 model 里名为 full_name（例如 "blocks.0.attn.qkv"）的子模块替换成 new_module。

    【Python 小知识】
        model.get_submodule("blocks.0.attn") 按点分隔的名字逐层找到子模块；
        setattr(parent, "qkv", new) 把父模块的这个属性换掉 —— nn.Module 会自动把新模块登记进去。
        对 nn.Sequential，子模块的名字是 "0"、"1"……，setattr(seq, "0", new) 同样有效。
    """
    parent_name, _, child_name = full_name.rpartition(".")
    parent = model.get_submodule(parent_name) if parent_name else model
    setattr(parent, child_name, new_module)


# =============================================================================
# 第 2 部分：LoRA
# =============================================================================
class LoRALinear(nn.Module):
    """
    LoRA（Low-Rank Adaptation，Hu 等 2021）：给一个冻结的线性层加一条低秩旁路。

        原来：     h = W x + b                    W 是 d_out × d_in，冻结
        加 LoRA： h = W x + b + (α / r) · B A x    A 是 r × d_in，B 是 d_out × r，只训练 A、B

    相当于把权重换成 W + ΔW，其中 ΔW = (α/r)·B A 的秩最多是 r（r 通常取 1~64，远小于 d）。

    【构造参数】
        base    : 要加 LoRA 的 nn.Linear（它的参数会被冻结）
        r       : 秩（rank），必须是正整数；越大越接近全量微调，参数也越多
        alpha   : 缩放系数 α；实际的缩放是 α / r（见下面的【思路】）
        dropout : 只作用在旁路输入上的 dropout（原论文的实现里有，默认不用）

    【可训练参数】
        A : (r, d_in)，随机初始化（高斯，标准差 1/√d_in）
        B : (d_out, r)，**初始化为 0**

    【思路】（讲义第 6 节）
        1. 为什么低秩？Aghajanyan 等（2020）发现预训练模型适配新任务时的"内在维度"很低；
           Hu 等进一步假设：微调带来的权重**变化量** ΔW 也是低秩的。所以只用 r(d_in + d_out) 个参数
           去表示 ΔW，而不是 d_in × d_out 个。
        2. 为什么 B 初始化为 0？这样训练开始时 ΔW = 0，模型和预训练的模型**完全一样**——
           微调从预训练模型出发，而不是从一个被随机扰动过的模型出发。
        3. 为什么 A 不能也初始化为 0？如果 A、B 都是 0，A 的梯度含有 B（= 0），B 的梯度含有 A x（= 0），
           两者的梯度都是 0，永远学不动（02 号文件演示 3 会验证）。所以必须有一个是随机的。
        4. 为什么要乘 α / r？换一个 r 时，不用重新调学习率（原论文：α 固定为第一次试的 r，之后不再调）。
        5. 训练完可以把 ΔW 加回 W（merge），推理时就只是一个普通的线性层，**没有任何额外的计算**。

    【边界】
        - merge() 之后再训练会出错（ΔW 被加了两次），所以 merged=True 时前向不再走旁路；要继续训练请先 unmerge()。
        - 如果 base 有偏置 b，LoRA 不改变它。
    """

    def __init__(self, base, r=4, alpha=8.0, dropout=0.0):
        super().__init__()
        if not isinstance(base, nn.Linear):
            raise TypeError("LoRALinear 只能包装 nn.Linear")
        if r <= 0:
            raise ValueError(f"r 必须是正整数，收到 {r}")
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        self.r, self.alpha = r, alpha
        self.scaling = alpha / r
        self.A = nn.Parameter(torch.randn(r, base.in_features) / math.sqrt(base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        self.drop = nn.Dropout(dropout)
        self.merged = False

    def delta_weight(self):
        """ΔW = (α/r)·B A，形状 (d_out, d_in)，和 base.weight 一样。"""
        return self.scaling * (self.B @ self.A)

    def forward(self, x):
        """
        【参数】x：(..., d_in)，前面可以有任意多维（比如 (B, T, d_in)）
        【返回】(..., d_out)

        【思路】先算 A x（降到 r 维），再算 B(A x)（升回 d_out 维）——
            计算量是 r(d_in + d_out)，而不是先算出 d_out × d_in 的 ΔW 再乘 x。
        """
        out = self.base(x)
        if not self.merged:
            out = out + self.scaling * (self.drop(x) @ self.A.t() @ self.B.t())
        return out

    @torch.no_grad()
    def merge(self):
        """把 ΔW 加进 base.weight。之后前向只走 base，结果不变、速度和普通线性层一样。"""
        if not self.merged:
            self.base.weight += self.delta_weight()
            self.merged = True

    @torch.no_grad()
    def unmerge(self):
        """merge() 的逆操作：从 base.weight 里减掉 ΔW（比如要换另一个任务的 LoRA 时）。"""
        if self.merged:
            self.base.weight -= self.delta_weight()
            self.merged = False


# 第 7 讲 TransformerLM 里各线性层的名字（"blocks.0.attn.qkv" 这样的名字以它们结尾）
ATTN_TARGETS = ("attn.qkv", "attn.proj")
ALL_TARGETS = ("attn.qkv", "attn.proj", "ffn.net.0", "ffn.net.2")


def add_lora(model, r=4, alpha=8.0, targets=ATTN_TARGETS, dropout=0.0):
    """
    把模型里名字以 targets 中某一项结尾的 nn.Linear 都换成 LoRALinear。

    【参数】
        model   : TransformerLM（或任何含 nn.Linear 的模型）
        r, alpha, dropout : 见 LoRALinear
        targets : 要加 LoRA 的层。默认只加在注意力的两个线性层上（qkv 和输出投影 proj）；
                  ALL_TARGETS 还包括前馈网络的两个线性层。

    【返回】被替换的层的名字列表

    【边界】
        - 不会冻结其他参数：通常先调用 freeze_all(model)，再 add_lora。
        - 我们的 qkv 是"一个 D → 3D 的线性层"（把 Wq、Wk、Wv 拼在一起），给它加 LoRA 相当于
          Wq、Wk、Wv 各有一个 B、但**共用同一个 A**。原论文是分别给 Wq、Wv 加，效果上差别不大。
        - 输出层（head）和词嵌入共用权重，这里不给它加 LoRA。

    【Python 小知识】
        list(model.named_modules()) 先把名字都收集起来再替换 —— 一边遍历一边修改模型结构会出问题。
    """
    names = [n for n, m in model.named_modules()
             if isinstance(m, nn.Linear) and any(n.endswith(t) for t in targets)]
    for n in names:
        _replace_module(model, n, LoRALinear(model.get_submodule(n), r=r, alpha=alpha, dropout=dropout))
    return names


def lora_layers(model):
    """返回模型里所有的 LoRALinear 层（名字 → 层）。"""
    return {n: m for n, m in model.named_modules() if isinstance(m, LoRALinear)}


def merge_lora(model):
    """把模型里所有 LoRA 层的 ΔW 合并进原来的权重。"""
    for m in lora_layers(model).values():
        m.merge()


def unmerge_lora(model):
    """撤销 merge_lora。"""
    for m in lora_layers(model).values():
        m.unmerge()


# =============================================================================
# 第 3 部分：适配器（Adapter）
# =============================================================================
class Adapter(nn.Module):
    """
    Houlsby 等（2019）的瓶颈适配器：  Adapter(h) = h + W_up · GELU(W_down · h)

        W_down : d → m（m 远小于 d，叫"瓶颈"）
        W_up   : m → d，**初始化为 0**，所以训练开始时 Adapter(h) = h（恒等映射）

    【构造参数】d：输入输出的维度；bottleneck：瓶颈宽度 m

    【思路】（讲义第 5.1 节）在每个子层（注意力、前馈网络）的输出后面插一个小网络，让它学"怎么修改这一层的输出"。
        自带的残差连接 + 零初始化保证插入之后模型的行为不变，训练从预训练模型出发。
        参数量约 2·d·m（每处），m = 8、d = 64 时约 1000 个。

    【边界】适配器是**串联**在主路径上的，推理时要多算几个小矩阵乘法（不像 LoRA 能合并进原权重）。
        对大模型的小批量推理，这一点额外延迟有时很明显（LoRA 论文表 1）。
    """

    def __init__(self, d, bottleneck=8):
        super().__init__()
        self.down = nn.Linear(d, bottleneck)
        self.up = nn.Linear(bottleneck, d)
        nn.init.normal_(self.down.weight, std=0.02)
        nn.init.zeros_(self.down.bias)
        nn.init.zeros_(self.up.weight)
        nn.init.zeros_(self.up.bias)

    def forward(self, h):
        return h + self.up(F.gelu(self.down(h)))


def add_adapters(model, bottleneck=8):
    """
    在 TransformerLM 每个块的两处插入适配器：注意力的输出投影之后、前馈网络之后（都在残差相加之前）。

    【返回】插入的适配器个数

    【思路】第 7 讲的块是  x = x + Drop(Attn(LN(x)))；  x = x + FFN(LN(x))。
        把 attn.proj 换成 Sequential(proj, Adapter)，把 ffn.net 换成 Sequential(net, Adapter)，
        适配器就恰好位于"子层输出之后、残差相加之前"——和 Houlsby 等的位置一样。
    """
    n = 0
    for block in model.blocks:
        d = block.attn.proj.out_features
        block.attn.proj = nn.Sequential(block.attn.proj, Adapter(d, bottleneck))
        block.ffn.net = nn.Sequential(block.ffn.net, Adapter(d, bottleneck))
        n += 2
    return n


# =============================================================================
# 第 4 部分：提示微调（Prompt Tuning）
# =============================================================================
class SoftPromptModel(nn.Module):
    """
    提示微调（Lester 等 2021）：在输入的词向量序列前面拼接 P 个**可训练的向量**（软提示），
    模型的所有参数冻结不动。

        原来的输入：  [<s>, w1, w2, ..., wn] 的词向量
        现在的输入：  [<s>, p1, p2, ..., pP, w1, ..., wn]    p1..pP 是可训练的 D 维向量

    【构造参数】
        lm        : 第 7 讲的 TransformerLM（会被冻结）
        n_prompt  : 软提示的长度 P
        init_ids  : None 或一组词编号 —— 用这些真实词的词向量初始化软提示（Lester 等发现这比随机初始化好）

    【思路】（讲义第 5.2 节）
        "提示工程"是人去找一串好用的**词**；但词是离散的，没法求梯度。
        提示微调把提示放宽成**连续的向量**，就可以用梯度下降去找"最好的提示"——
        它不必对应任何真实的词。可训练参数只有 P × D 个（P = 8、D = 64 时只有 512 个）。

    【边界】
        - 序列变长了 P 个位置：输出的 logits / hidden 长度是 T + P，使用时要把下标往后挪 P（见 forward 的说明）。
        - 总长度 T + P 不能超过模型的 max_len（可学习的位置嵌入只有那么多行）。
        - 只在模型**输入**处加提示（Lester 等的做法）；前缀微调（Li & Liang 2021）在**每一层**的
          键和值前面都加可训练的向量，参数多一些、在小模型上更稳定（讲义第 5.2 节、练习 5）。
    """

    def __init__(self, lm, n_prompt=8, init_ids=None):
        super().__init__()
        self.lm = freeze_all(lm)
        self.n_prompt = n_prompt
        emb = lm.tok_emb.weight
        if init_ids is not None:
            idx = torch.tensor([init_ids[i % len(init_ids)] for i in range(n_prompt)])
            init = emb[idx].detach().clone()
        else:
            init = torch.randn(n_prompt, emb.shape[1]) * emb.detach().std()
        self.prompt = nn.Parameter(init)

    def forward(self, X, pad_mask=None):
        """
        【参数】
            X        : (B, T) 词编号，第 0 个位置必须是 <s>
            pad_mask : None 或 (B, T) 布尔张量（True = 真实的词）

        【返回】(logits, hidden)，形状 (B, T + P, V) 和 (B, T + P, D)。
            原序列第 t 个位置（t ≥ 1）在输出里的下标是 t + P；<s> 仍在下标 0。

        【思路】重写 TransformerLM.forward 的前半部分：先查词向量，插入软提示，再加位置嵌入，
            然后照常经过所有块。（我们不能直接调用 lm(X)，因为它只接受词编号，不接受向量。）
        """
        lm = self.lm
        B, T = X.shape
        tok = lm.tok_emb(X)                                               # (B, T, D)
        P = self.prompt.unsqueeze(0).expand(B, -1, -1)                    # (B, P, D)
        h = torch.cat([tok[:, :1], P, tok[:, 1:]], dim=1)                 # (B, T + P, D)
        L = T + self.n_prompt
        if L > lm.max_len:
            raise ValueError(f"加上软提示后长度 {L} 超过了 max_len={lm.max_len}")
        if lm.pos == "learned":
            h = h + lm.pos_emb(torch.arange(L, device=X.device))
        elif lm.pos == "sinusoidal":
            h = h + lm.pos_table[:L]
        h = lm.drop(h)
        if pad_mask is not None:
            ones = torch.ones(B, self.n_prompt, dtype=torch.bool, device=X.device)
            pad_mask = torch.cat([pad_mask[:, :1], ones, pad_mask[:, 1:]], dim=1)
        for block in lm.blocks:
            h = block(h, pad_mask)
        hidden = lm.ln_f(h)
        return lm.head(hidden), hidden


# =============================================================================
# 自我检查
# =============================================================================
if __name__ == "__main__":
    import sys
    from pathlib import Path

    sys.path.append(str(Path(__file__).resolve().parent.parent / "lecture07"))
    from transformer import TransformerLM

    torch.manual_seed(0)
    X = torch.randint(1, 50, (2, 10))

    # 1. LoRA：B = 0 时输出和原模型完全一样；随机 B 时，merge 前后输出一样
    lm = TransformerLM(50, d_model=32, n_heads=4, n_layers=2, dropout=0.0).eval()
    ref = lm(X)[0]
    freeze_all(lm)
    names = add_lora(lm, r=2, alpha=4)
    assert torch.allclose(lm(X)[0], ref, atol=1e-6)
    for m in lora_layers(lm).values():
        nn.init.normal_(m.B, std=0.1)
    before = lm(X)[0]
    assert not torch.allclose(before, ref, atol=1e-3)
    merge_lora(lm)
    assert torch.allclose(lm(X)[0], before, atol=1e-5)
    unmerge_lora(lm)
    assert torch.allclose(lm(X)[0], before, atol=1e-5)
    n_tr, n_all = count_trainable(lm)
    print(f"  LoRA：替换了 {len(names)} 个线性层，可训练参数 {n_tr:,} / {n_all:,}；初始化、合并、撤销合并都检查通过")

    # 2. 适配器：插入后（零初始化）输出不变
    lm = TransformerLM(50, d_model=32, n_heads=4, n_layers=2, dropout=0.0).eval()
    ref = lm(X)[0]
    freeze_all(lm)
    k = add_adapters(lm, bottleneck=4)
    assert torch.allclose(lm(X)[0], ref, atol=1e-6)
    print(f"  适配器：插入了 {k} 个，可训练参数 {count_trainable(lm)[0]:,}；插入后输出不变")

    # 3. BitFit
    lm = TransformerLM(50, d_model=32, n_heads=4, n_layers=2, dropout=0.0)
    bitfit(lm)
    print(f"  BitFit：可训练参数 {count_trainable(lm)[0]:,} / {count_trainable(lm)[1]:,}")

    # 4. 提示微调：形状正确，只有软提示可训练
    lm = TransformerLM(50, d_model=32, n_heads=4, n_layers=2, dropout=0.0, max_len=32).eval()
    sp = SoftPromptModel(lm, n_prompt=5, init_ids=[7, 8, 9])
    logits, hidden = sp(X)
    assert logits.shape == (2, 15, 50) and hidden.shape == (2, 15, 32)
    assert count_trainable(sp)[0] == 5 * 32
    print(f"  提示微调：输出长度 10 + 5 = {logits.shape[1]}，可训练参数 {count_trainable(sp)[0]}")
    print("  全部自我检查通过 ✓")