"""
peft.py —— 第 6 讲：从零实现参数高效微调（PEFT）的几种主要方法
=====================================================================

【这个文件做什么】
    把第 2 讲的 TransformerLM（d=64、2 层、4 头、约 10.7 万参数）"改装"成各种 PEFT 版本。
    按 Delta Tuning（Ding et al. 2022）和第 6 讲 slides 的分法，三大类：

    ① 加法（Addition）：加入新的小参数，只训练它们
        SoftPromptLM      prompt tuning（Lester et al. 2021）：输入前面接 P 个可学习的"软 token"嵌入
        add_prefix        prefix tuning（Li & Liang 2021）：每一层注意力的 K、V 前面都接 P 个可学习向量
        add_adapters      adapter（Houlsby et al. 2019）：每个子层输出后面串一个"瓶颈 MLP"，初始化成恒等
    ② 指定（Specification）：只训练原模型的一小部分参数
        bitfit_params     BitFit（Ben Zaken et al. 2022）：只训练偏置项
        last_layer_params 层冻结：只训练最后一层
    ③ 重参数化（Reparameterization）：把"权重的改变量"写成低维的形式
        add_lora          LoRA（Hu et al. 2021）：ΔW = B A，秩 r ≪ d
        add_ia3           (IA)³（Liu et al. 2022）：用可学习的向量逐元素缩放 K、V 和 FFN 中间激活

    每个 add_* 函数都**就地**修改模型（冻结原参数），返回新加的可训练参数列表；
    模型的 forward 接口不变（forward(ids, pad=...) → logits），所以训练、生成、评估代码完全通用。

【为什么不直接用 HuggingFace PEFT 库】
    库把一切都藏起来了。自己写一遍，每种方法"改了哪里、加了多少参数、初始化为什么是那样"一目了然；
    写完再看 PEFT 库的 LoraConfig(r=..., target_modules=["q_proj", "v_proj"]) 就知道它在做什么。

【LLM 联系】
    这些方法在 7B～70B 的模型上原理完全一样，区别只在规模：Llama-2-7B 全参数微调要更新 67 亿个参数，
    r=8 的 LoRA（只加在 q、v 上）只训练约 400 万个——0.06%。
"""

import math

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

from transformer import MultiHeadAttention, scaled_dot_product_attention


def freeze(model):
    """冻结模型的全部参数（PEFT 的第一步）。"""
    for p in model.parameters():
        p.requires_grad_(False)


# =====================================================================
# ③ 重参数化：LoRA
# =====================================================================

class LoRALinear(nn.Module):
    """
    y = W x + b + (α / r) · B A x

    【思路】
        全参数微调学的是 W' = W + ΔW。Aghajanyan et al. (2021) 发现微调的"内在维度"很低——
        在一个很小的随机子空间里优化就能达到全参数微调 90% 的效果。LoRA 更进一步假设：
        **ΔW 本身是低秩的**，于是把它写成 B (d_out × r) 乘 A (r × d_in)，只训练 A 和 B。
        参数量从 d_out·d_in 降到 r·(d_in + d_out)。
    【初始化】A 用随机数，B 用全 0 → 训练开始时 ΔW = 0，模型和原来一模一样
             （这一点很重要：从"已经很好"的起点出发，而不是一开始就把模型搞乱）。
             如果 A、B 都是 0，梯度也全是 0，永远学不动；都随机，起点就被破坏了。
    【参数】base：原来的 nn.Linear（被冻结）；r：秩；alpha：缩放系数（α / r 让"换秩时不用重调学习率"）。
    【边界】merge() 把 BA 加回 W，之后推理和原模型一样快——LoRA 相对 adapter 的一大优点（04 号脚本）。
    """

    def __init__(self, base, r=4, alpha=None):
        super().__init__()
        self.base = base
        self.r = r
        self.scale = (alpha if alpha is not None else r) / r
        self.A = nn.Parameter(torch.empty(r, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))      # 和 nn.Linear、HF PEFT 的默认初始化相同
        self.merged = False

    def forward(self, x):
        y = self.base(x)
        if not self.merged:
            y = y + (x @ self.A.T @ self.B.T) * self.scale
        return y

    def delta(self):
        """ΔW = (α/r)·B A，形状和 base.weight 相同。"""
        return (self.B @ self.A) * self.scale

    @torch.no_grad()
    def merge(self):
        """把 ΔW 并进原权重（之后 forward 不再多算两次矩阵乘法）。"""
        if not self.merged:
            self.base.weight += self.delta()
            self.merged = True


def add_lora(model, r=4, alpha=None, targets=("q_proj", "v_proj")):
    """
    给模型里名字以 targets 结尾的线性层装上 LoRA。

    【参数】targets：默认 ("q_proj", "v_proj")——LoRA 论文的默认选择（只改注意力的 Q、V）；
            可以加上 "k_proj"、"o_proj"、"up"、"down"（FFN 的两层）。
    【返回】新加的参数列表 [A1, B1, A2, B2, ...]。
    【Python 小知识】named_modules() 遍历所有子模块及其"路径名"（例如 blocks.0.self_attn.q_proj）；
            要替换一个子模块，得先找到它的父模块，再 setattr(父, 属性名, 新模块)。
    """
    freeze(model)
    params = []
    for name, mod in list(model.named_modules()):
        for child_name, child in list(mod.named_children()):
            if isinstance(child, nn.Linear) and child_name in targets:
                lora = LoRALinear(child, r=r, alpha=alpha)
                setattr(mod, child_name, lora)
                params += [lora.A, lora.B]
    return params


def lora_modules(model):
    return [m for m in model.modules() if isinstance(m, LoRALinear)]


# =====================================================================
# ③ 重参数化：(IA)³
# =====================================================================

class ScaleOut(nn.Module):
    """y = l ⊙ f(x)：对一个线性层的输出逐元素乘一个可学习向量 l（初始化为全 1 = 不改变）。"""

    def __init__(self, base):
        super().__init__()
        self.base = base
        self.l = nn.Parameter(torch.ones(base.out_features))

    def forward(self, x):
        return self.base(x) * self.l


class ScaleIn(nn.Module):
    """y = f(l ⊙ x)：对一个线性层的**输入**逐元素缩放。用在 FFN 的 down 层上 = 缩放 FFN 的中间激活。"""

    def __init__(self, base):
        super().__init__()
        self.base = base
        self.l = nn.Parameter(torch.ones(base.in_features))

    def forward(self, x):
        return self.base(x * self.l)


def add_ia3(model):
    """
    (IA)³ = Infused Adapter by Inhibiting and Amplifying Inner Activations。

    【思路】不学"加多少"，只学"放大 / 抑制多少"：每层三个向量 l_k、l_v（d 维）、l_ff（d_ff 维），
            分别逐元素乘在 K、V 和 FFN 中间激活上。初始化为 1，训练开始时模型不变。
            参数极少（这里每层 64+64+256 = 384 个），而且因为只是逐元素乘法，
            **同一个批次里不同样本可以用不同任务的向量**（slides 说的 mixed-task batches）。
    【返回】新加的参数列表。
    """
    freeze(model)
    params = []
    for blk in model.blocks:
        att = blk.self_attn
        att.k_proj = ScaleOut(att.k_proj)
        att.v_proj = ScaleOut(att.v_proj)
        blk.ffn.down = ScaleIn(blk.ffn.down)
        params += [att.k_proj.l, att.v_proj.l, blk.ffn.down.l]
    return params


# =====================================================================
# ① 加法：adapter
# =====================================================================

class Adapter(nn.Module):
    """
    瓶颈 MLP：h ← h + W_up · GELU(W_down · h)，d → m → d，m ≪ d。

    【初始化】W_up 全 0 → 刚装上时整个 adapter 是恒等映射（slides："初始化成和原模型输出一样"）。
    """

    def __init__(self, d, m):
        super().__init__()
        self.down = nn.Linear(d, m)
        self.up = nn.Linear(m, d)
        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)))


class WithAdapter(nn.Module):
    """把一个子层（注意力或 FFN）包起来：先算原子层，再过 adapter。注意力返回 (out, weights)，要分开处理。"""

    def __init__(self, sub, d, m):
        super().__init__()
        self.sub = sub
        self.adapter = Adapter(d, m)

    def forward(self, *args, **kwargs):
        out = self.sub(*args, **kwargs)
        if isinstance(out, tuple):
            return (self.adapter(out[0]),) + tuple(out[1:])
        return self.adapter(out)


def add_adapters(model, m=8):
    """
    Houlsby 式 adapter：每个块的注意力子层和 FFN 子层后面各插一个（在残差相加之前）。

    【参数】m：瓶颈宽度。每个 adapter 有 2·d·m + m + d 个参数。
    【返回】新加的参数列表。
    【边界】adapter 是**串联**在计算路径上的额外层：推理时也要多算这几层，没法像 LoRA 那样"合并"掉。
    """
    freeze(model)
    params = []
    d = model.d
    for blk in model.blocks:
        blk.self_attn = WithAdapter(blk.self_attn, d, m)
        blk.ffn = WithAdapter(blk.ffn, d, m)
        params += list(blk.self_attn.adapter.parameters()) + list(blk.ffn.adapter.parameters())
    return params


# =====================================================================
# ① 加法：prefix tuning
# =====================================================================

class PrefixAttention(nn.Module):
    """
    在 K、V 的最前面接上 P 个可学习的向量：每个 query 除了看真实 token，还能看这 P 个"虚拟 token"。

    【思路】prompt tuning 只在**输入层**加软 token，要靠后面的层把它的影响"传上去"；
            prefix tuning 直接在**每一层**提供可学习的 K、V，影响力更直接、参数也更多（每层 2·P·d）。
    【参数】base：原来的 MultiHeadAttention（冻结）；P：前缀长度；init_kv：可选 (P, d)、(P, d) 的初始值。
    【边界】掩码：原掩码形状能广播到 (B, h, T, T)；前缀对所有 query 都可见，所以在左边补 P 列 True。
    """

    def __init__(self, base, P, init_kv=None):
        super().__init__()
        self.base = base
        d = base.h * base.hd
        if init_kv is None:
            self.pk = nn.Parameter(torch.randn(P, d) * 0.5)
            self.pv = nn.Parameter(torch.randn(P, d) * 0.5)
        else:
            self.pk = nn.Parameter(init_kv[0].clone())
            self.pv = nn.Parameter(init_kv[1].clone())

    def forward(self, x, kv=None, mask=None):
        b = self.base
        B, T, _ = x.shape
        P = self.pk.shape[0]
        q = b._split(b.q_proj(x))
        k = b._split(torch.cat([self.pk.expand(B, P, -1), b.k_proj(x)], 1))
        v = b._split(torch.cat([self.pv.expand(B, P, -1), b.v_proj(x)], 1))
        if mask is not None:
            mask = mask.expand(B, 1, T, T)       # (T, T) 或 (B, 1, T, T) 都能这样扩展
            mask = torch.cat([torch.ones(B, 1, T, P, dtype=torch.bool), mask], -1)
        out, w = scaled_dot_product_attention(q, k, v, mask)
        out = out.transpose(1, 2).reshape(B, T, -1)
        return b.o_proj(out), w[..., P:]


@torch.no_grad()
def real_kv(model, token_ids):
    """
    用真实 token 跑一遍模型，取每一层注意力的 K、V 投影（投影之后、拆头之前），作为前缀的初始值。

    【返回】[(K (P, d), V (P, d)) 每层一个]。
    【为什么】Li & Liang 发现用真实词的激活初始化前缀，比随机初始化好得多（slides "How to initialize the prefix?"）：
            随机向量落在模型从没见过的区域，训练要先花很多步"把它拉回来"。
    """
    ids = torch.tensor([token_ids])
    x = model.embed(ids)
    T = ids.shape[1]
    mask = model.make_mask(T)
    out = []
    for blk in model.blocks:
        h = blk.n1(x)
        out.append((blk.self_attn.k_proj(h)[0], blk.self_attn.v_proj(h)[0]))
        x, _ = blk(x, mask=mask)
    return out


def add_prefix(model, P=8, init_ids=None):
    """
    给每一层的自注意力装上长度 P 的前缀。

    【参数】init_ids：可选，长度 P 的 token id 列表；给了就用这些真实 token 的 K、V 初始化（否则随机）。
    【返回】新加的参数列表（每层 pk、pv）。
    """
    freeze(model)
    init = real_kv(model, init_ids) if init_ids is not None else [None] * len(model.blocks)
    params = []
    for blk, kv in zip(model.blocks, init):
        blk.self_attn = PrefixAttention(blk.self_attn, P, kv)
        params += [blk.self_attn.pk, blk.self_attn.pv]
    return params


# =====================================================================
# ① 加法：prompt tuning
# =====================================================================

class SoftPromptLM(nn.Module):
    """
    prompt tuning：在输入 token 的嵌入前面接 P 个可学习的向量，其余一切冻结。

    【思路】人写提示（prompt engineering）要反复试词语；既然有任务数据，为什么不直接**用梯度下降找提示**？
            离散的 token 没法求导，于是退一步：优化连续的嵌入向量（"软提示"）。
            它们不必对应任何真实的词——这也是 slides 说"学到的嵌入不好解释"的原因。
    【参数】base：TransformerLM（冻结）；P：软提示长度；init_ids：可选，用这些真实 token 的嵌入初始化
            （Lester et al. 发现用词表里的词、尤其是类别标签词初始化，比随机初始化好）。
    【前向】ids (B, T) → logits (B, T, V)：内部序列是 [软提示 P 个, ids 的 T 个]，输出时去掉前 P 个位置，
            所以对调用者来说接口和原模型完全一样。
    【边界】软提示占用序列长度：每个前缀位置都是从真实任务"借"来的上下文（slides 的 pitfall）。
            位置编码也跟着右移了 P 格。
    """

    def __init__(self, base, P=8, init_ids=None):
        super().__init__()
        self.base = base
        freeze(base)
        d = base.d
        if init_ids is not None:
            init = base.emb.weight[torch.tensor(init_ids)].detach().clone()
        else:
            init = torch.randn(P, d) * d ** -0.5            # 和词嵌入的初始化同一个尺度
        self.prompt = nn.Parameter(init)
        self.d = d

    def forward(self, ids, pad=None):
        b = self.base
        B, T = ids.shape
        P = self.prompt.shape[0]
        tok = b.emb(ids)
        x = torch.cat([self.prompt.expand(B, P, -1), tok], 1) * math.sqrt(self.d)
        # 位置编码按总长度 P + T 计算（复用第 2 讲 TransformerLM.embed 的逻辑）
        from transformer import sinusoidal_positions
        x = x + sinusoidal_positions(P + T, self.d)
        pad_full = None if pad is None else torch.cat([torch.zeros(B, P, dtype=torch.bool), pad], 1)
        mask = b.make_mask(P + T, pad_full)
        for blk in b.blocks:
            x, _ = blk(x, mask=mask)
        x = b.norm_f(x)[:, P:]
        return x @ b.emb.weight.T


# =====================================================================
# ② 指定：BitFit 与层冻结
# =====================================================================

def bitfit_params(model):
    """
    BitFit：只训练所有的偏置项（线性层的 bias、LayerNorm 的 bias）。

    【思路】偏置只是给每个神经元的输入"挪一挪零点"，参数量不到 1.5%，
            但 Ben Zaken et al. 发现在 BERT 上它能做到接近全参数微调的效果——说明微调很大程度上是
            "调整已有特征的开关和阈值"，而不是学新特征。
    【返回】参数列表（没有冻结的工作：finetune() 会只对这些参数打开梯度）。
    """
    return [p for n, p in model.named_parameters() if n.endswith("bias")]


def last_layer_params(model):
    """层冻结：只训练最后一个 Transformer 块和最后的 LayerNorm（以及它们的偏置）。"""
    return list(model.blocks[-1].parameters()) + list(model.norm_f.parameters())


def all_params(model):
    return list(model.parameters())


# =====================================================================
# 量化（QLoRA 的思路，04 号脚本用）
# =====================================================================

@torch.no_grad()
def quantize_(model, bits=4, block=32):
    """
    把模型里所有线性层和嵌入的权重"量化再反量化"（模拟只存 bits 位整数的效果）。

    【思路】分块绝对值最大（absmax）量化：每 block 个权重一组，按组内最大绝对值 m 缩放到
            [-(2^(bits-1)-1), 2^(bits-1)-1] 的整数，再乘回去。误差最大约 m / (2^(bits-1)-1) / 2。
            QLoRA 用的是更讲究的 4 位 NormalFloat（按正态分布的分位数取 16 个值）和双重量化，
            思路相同：**冻结的基础模型只需要存成 4 位**，LoRA 的 A、B 仍然用高精度训练。
    【返回】None（就地修改）。
    【边界】只模拟了"精度损失"，没有真的省内存——真正的 4 位存储要专门的 CUDA 内核（bitsandbytes）。
    """
    qmax = 2 ** (bits - 1) - 1
    for mod in model.modules():
        W = None
        if isinstance(mod, nn.Linear):
            W = mod.weight
        elif isinstance(mod, nn.Embedding):
            W = mod.weight
        if W is None:
            continue
        flat = W.reshape(-1)
        n = flat.numel()
        padn = (-n) % block
        g = torch.cat([flat, flat.new_zeros(padn)]).view(-1, block)
        m = g.abs().amax(1, keepdim=True).clamp_min(1e-12)
        q = torch.round(g / m * qmax).clamp(-qmax, qmax)
        W.copy_((q * m / qmax).reshape(-1)[:n].view_as(W))


__all__ = ["freeze", "LoRALinear", "add_lora", "lora_modules", "add_ia3", "add_adapters", "add_prefix",
           "SoftPromptLM", "bitfit_params", "last_layer_params", "all_params", "quantize_", "real_kv",
           "MultiHeadAttention"]