"""Granite-1B-A400M 的 Looped × AttnRes 改造(消融实验 G0-G3 共用一份代码)。

开关(通过 from_pretrained 的 kwargs 直接传给 __init__):
- loop_rounds: int = 1      循环轮数; 1 = G0 基线
- attnres: bool = False     轮内 Block AttnRes (G2+)
- cross_round: bool = False 源集合加入之前所有轮次的块摘要 (G3)
- attnres_blocks: int = 8   AttnRes 块数 (48 子层 / 8 块 = 6 子层每块)

结构要点(依据 DESIGN.md 与两篇论文):
- 循环: 24 层 × R 轮, 权重跨轮硬共享; 每轮输出过共享 lm_head 算 CE, 取各轮均值
- AttnRes: 每子层(attn/MoE 各算一个)一个零初始化可学习伪查询, 按轮独立;
  key=value=子层输出, 核 exp(q·RMSNorm(k)) (RMSNorm 无可学习参数), softmax 在源维度
- residual_multiplier μ=0.22 折进 value (v_i = μ·f_i)
- 块内: 部分和按标准残差累加; 块边界: 块摘要加入源集合
- 每轮末尾: 用独立的终读出伪查询对 [b0] + (跨轮时含历史块) + 本轮块摘要 做最后读出
- 跨轮(G3): 源 = [b0(embedding)] + [之前所有轮次的块摘要] + [本轮已完成块] + [本块部分和]

注意: 训练前向 (input_ids/labels) 不支持 KV cache; 推理用 greedy_generate(手动贪心 + KV cache)。
"""
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.checkpoint import checkpoint
from types import SimpleNamespace

from transformers.models.granitemoe.configuration_granitemoe import GraniteMoeConfig
from transformers.models.granitemoe.modeling_granitemoe import (
    GraniteMoeForCausalLM,
    load_balancing_loss_func,
)
from transformers.masking_utils import create_causal_mask


def _rms_norm(x, eps=1e-6):
    """无可学习参数的 RMSNorm(AttnRes 核里用在 key 上)。"""
    return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps).to(x.dtype)


class GraniteLoopForCausalLM(GraniteMoeForCausalLM):
    def __init__(self, config: GraniteMoeConfig, loop_rounds: int = 1, attnres: bool = False,
                 cross_round: bool = False, attnres_blocks: int = 8, **kwargs):
        super().__init__(config)
        # 开关: kwargs 优先, 否则读 config(兼容加载我们自己 save_pretrained 的 checkpoint)
        self.loop_rounds = getattr(config, "loop_rounds", None) or loop_rounds
        self.use_attnres = getattr(config, "attnres", None)
        if self.use_attnres is None:
            self.use_attnres = attnres
        self.cross_round = getattr(config, "cross_round", None)
        if self.cross_round is None:
            self.cross_round = cross_round
        self.attnres_blocks = getattr(config, "attnres_blocks", None) or attnres_blocks
        # 写回 config, 使 save_pretrained 能持久化开关
        config.loop_rounds = self.loop_rounds
        config.attnres = self.use_attnres
        config.cross_round = self.cross_round
        config.attnres_blocks = self.attnres_blocks
        L = config.num_hidden_layers
        assert (2 * L) % self.attnres_blocks == 0, "子层数需被块数整除"
        self.sub_per_block = (2 * L) // self.attnres_blocks  # 每块子层数
        assert self.sub_per_block % 2 == 0
        self.layers_per_block = self.sub_per_block // 2

        d = config.hidden_size
        self.alpha_recorder = None  # 分析钩子: list 时 _attnres_step 会记录各源注意力权重
        self.round_recorder = None  # 分析钩子: list 时 forward 会记录每轮结束后的 hidden
        if self.use_attnres:
            # 伪查询: (轮数, 子层数 + 1, d), 零初始化 -> 初始为等权平均
            # +1 = 每轮结束时的终读出查询(对 [b0] + 各块摘要 做最后一次注意力,
            # 否则最终 hidden 只剩最后一个块的部分和, 丢失 b0 主干 -> 初始 loss 崩到随机水平)
            self.pseudo_queries = nn.Parameter(
                torch.zeros(self.loop_rounds, 2 * L + 1, d)
            )

    @torch.no_grad()
    def reset_pseudo_queries(self):
        """from_pretrained 后必须调用: transformers 5 的加载管线会把 checkpoint 里
        缺失的 pseudo_queries 重置为未初始化内存(可能含 NaN), 这里重新归零,
        恢复论文要求的零初始化(初始 = 等权平均)。自己 save_pretrained 的
        checkpoint 含训练后的值, 加载后无需调用。"""
        if self.use_attnres:
            self.pseudo_queries.zero_()

    # ---------- 单个子层前向 ----------
    def _sublayer(self, layer, kind, h, mask, pos_emb):
        if kind == "attn":
            return layer.self_attn(
                hidden_states=layer.input_layernorm(h),
                attention_mask=mask,
                position_embeddings=pos_emb,
            )[0]
        # moe: 复制 GraniteMoeMoE.forward 但保留 router_logits(aux loss 用),
        # 作为 checkpoint 段的返回值传出, 避免用 hook 持有段内张量引用(会阻止激活释放)
        moe = layer.block_sparse_moe
        normed = layer.post_attention_layernorm(h)
        b, t, d = normed.shape
        flat = normed.reshape(-1, d)
        top_idx, top_w, rlogits = moe.router(flat)
        out = moe.experts(flat, top_idx, top_w).view(b, t, d)
        return out, rlogits

    def _aggregate(self, q, sources):
        """AttnRes 聚合: q (d,), sources list of (b,s,d) -> (b,s,d)"""
        V = torch.stack(sources)                       # (S,b,t,d)
        logits = torch.einsum("d,Sbtd->Sbt", q, _rms_norm(V))
        alpha = logits.softmax(dim=0)                  # 源维度归一化
        return torch.einsum("Sbt,Sbtd->btd", alpha, V)

    # ---------- checkpoint 段(显存优化: 粒度从整轮细化到单层/单块) ----------
    # 整轮一个段时, 反向重算会把 24 层(AttnRes 时还有每子层 stack 出的全部源副本 V)
    # 同时留在显存里 -> OOM。细化后每段反向只重算一层/一个块。
    def _layer_fwd(self, layer, h, mask, pos_emb, mu):
        """非 AttnRes 路径: 一个完整层(attn + moe), 返回 (h, router_logits)。"""
        f = self._sublayer(layer, "attn", h, mask, pos_emb)
        h = h + mu * f
        f, rl = self._sublayer(layer, "moe", h, mask, pos_emb)
        return h + mu * f, rl

    def _attnres_step(self, layer, kind, q, cat, inv_rms, block_sum, mask, pos_emb, mu):
        """AttnRes 路径: 聚合 + 单个子层 + 块内累加, 返回 (新 block_sum, router_logits 或 None)。

        性能优化版聚合(profile 显示旧的"每子层 stack 全部源 + 对全部源做 RMSNorm"
        贡献了 ~60% 的 CUDA 时间, 全是 elementwise/cat):
        - cat: (S,b,t,d) 块内固定源的堆叠, 每块只 stack 一次(块级缓存)
        - inv_rms: (S,b,t) 各固定源的 RMSNorm 系数, 利用 q·RMSNorm(k) = (q·k)·inv_rms
        - softmax 拆成 [固定源 | block_sum] 两段, 每子层只需对 block_sum 单张量求范数
        数学上与 _aggregate 完全等价(仅浮点运算次序不同)。"""
        logits = torch.einsum("d,sbtd->sbt", q, cat) * inv_rms
        if block_sum is not None:
            bs_inv = torch.rsqrt(block_sum.float().pow(2).mean(-1) + 1e-6)
            bs_logit = torch.einsum("d,btd->bt", q, block_sum) * bs_inv
            alpha = torch.cat([logits, bs_logit.unsqueeze(0)], dim=0).softmax(dim=0)
        else:
            alpha = logits.softmax(dim=0)
        # 分析钩子: 仅分析脚本设置 alpha_recorder 时生效(训练时恒为 None, 零开销)。
        # 记录 (上下文(轮,块,层序,类型), 各源平均权重), 用于检验"跨轮检索是否偷懒"。
        rec = getattr(self, "alpha_recorder", None)
        if rec is not None:
            with torch.no_grad():
                rec.append((self._rec_ctx, alpha.detach().mean(dim=(1, 2)).cpu()))
        alpha = alpha.to(cat.dtype)
        if block_sum is not None:
            h = torch.einsum("sbt,sbtd->btd", alpha[:-1], cat) + alpha[-1].unsqueeze(-1) * block_sum
        else:
            h = torch.einsum("sbt,sbtd->btd", alpha, cat)
        f = self._sublayer(layer, kind, h, mask, pos_emb)
        rl = None
        if kind == "moe":
            f, rl = f
        upd = mu * f
        return (upd if block_sum is None else block_sum + upd), rl

    def _run_block(self, blk_layers, q_blk, fixed, mask, pos_emb, mu):
        """一个 AttnRes 块(3 层 / 6 个子层), 作为一个 checkpoint 段。

        源堆叠 cat 与 inv_rms 在段内构建(而非段外传入): 段内张量在前向结束后即释放,
        反向重算时重建, 避免 32 个块的 cat(~0.5GB each)同时常驻显存。
        profile 显示旧的"每子层 stack 全部源 + 对全部源 RMSNorm"贡献 ~60% CUDA 时间,
        这里每块只 stack/归一化一次, 块内 6 个子层复用, 并用
        q·RMSNorm(k) = (q·k)·inv_rms 把归一化折成标量乘法(数学上与 _aggregate 等价)。"""
        cat = torch.stack(fixed)                                      # (S,b,t,d)
        inv_rms = torch.rsqrt(cat.float().pow(2).mean(-1) + 1e-6)     # (S,b,t) fp32
        block_sum = None
        rlogits = []
        for j, layer in enumerate(blk_layers):
            for kind in ("attn", "moe"):
                q = q_blk[2 * j] if kind == "attn" else q_blk[2 * j + 1]
                self._rec_ctx = (*self._rec_ctx_base, j, kind)  # 分析钩子用, 训练时无读取方
                block_sum, rl = self._attnres_step(
                    layer, kind, q, cat, inv_rms, block_sum, mask, pos_emb, mu)
                if rl is not None:
                    rlogits.append(rl)
        return block_sum, rlogits

    def _ce(self, h, labels):
        """单轮 CE, 单独成段: 避免 4 轮的 logits(fp32 上转后每轮 ~1.2GB)常驻反传图。"""
        logits = self.lm_head(self.model.norm(h)) / self.logits_scaling
        return self.loss_function(logits, labels, vocab_size=self.config.vocab_size)

    # ---------- 一轮循环 ----------
    def _run_round(self, h, round_idx, b0, prev_blocks, mask, pos_emb, mu, use_ckpt):
        """返回 (本轮最终 hidden, 块摘要列表, 各 MoE 层的 router_logits)。"""
        layers = self.model.layers
        rlogits = []
        if not self.use_attnres:
            for layer in layers:
                if use_ckpt:
                    h, rl = checkpoint(
                        lambda hh, l=layer: self._layer_fwd(l, hh, mask, pos_emb, mu),
                        h, use_reentrant=False)
                else:
                    h, rl = self._layer_fwd(layer, h, mask, pos_emb, mu)
                rlogits.append(rl)
            block_sums = []
        else:
            fixed = [b0] + (list(prev_blocks) if self.cross_round else [])
            block_sums = []
            for blk in range(self.attnres_blocks):
                blk_layers = layers[blk * self.layers_per_block:(blk + 1) * self.layers_per_block]
                q_blk = self.pseudo_queries[round_idx,
                                            blk * self.sub_per_block:(blk + 1) * self.sub_per_block]
                fix_snap = tuple(fixed)  # fixed 之后会 append, 必须快照, 否则重算时形状对不上
                self._rec_ctx_base = (round_idx, blk)  # 分析钩子用
                if use_ckpt:
                    block_sum, rls = checkpoint(
                        lambda _q, bl=blk_layers, qb=q_blk, fs=fix_snap:
                            self._run_block(bl, qb, fs, mask, pos_emb, mu),
                        q_blk, use_reentrant=False)
                else:
                    block_sum, rls = self._run_block(blk_layers, q_blk, fix_snap, mask, pos_emb, mu)
                rlogits.extend(rls)
                fixed.append(block_sum)
                block_sums.append(block_sum)
            # 终读出: 对 [b0] + (跨轮时的历史块) + 本轮各块摘要 做最后一次注意力,
            # 而不是只取最后一个块的部分和
            final_sources = [b0] + (list(prev_blocks) if self.cross_round else []) + block_sums
            q_final = self.pseudo_queries[round_idx, 2 * len(layers)]
            if use_ckpt:
                h = checkpoint(
                    lambda _q, qq=q_final, src=tuple(final_sources): self._aggregate(qq, list(src)),
                    q_final, use_reentrant=False)
            else:
                h = self._aggregate(q_final, final_sources)
        return h, block_sums, rlogits

    # ---------- 训练前向 ----------
    def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
        model = self.model
        bsz, seq = input_ids.shape
        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0)
        h = model.embed_tokens(input_ids) * model.embedding_multiplier
        mask = create_causal_mask(
            config=self.config, inputs_embeds=h, attention_mask=attention_mask,
            past_key_values=None, position_ids=position_ids,
        )
        pos_emb = model.rotary_emb(h, position_ids)
        mu = self.config.residual_multiplier
        b0 = h

        use_ckpt = self.training and torch.is_grad_enabled()
        ce_losses = []
        all_rlogits = []
        prev_blocks = []  # cross_round 时跨轮累积的块摘要
        for r in range(self.loop_rounds):
            h, round_blocks, rlogits = self._run_round(
                h, r, b0, prev_blocks, mask, pos_emb, mu, use_ckpt)
            if self.round_recorder is not None:
                self.round_recorder.append(h.detach())
            if self.cross_round:
                prev_blocks = prev_blocks + round_blocks
            if labels is not None:
                if use_ckpt:
                    ce = checkpoint(lambda hh=h, lb=labels: self._ce(hh, lb),
                                    h, use_reentrant=False)
                else:
                    ce = self._ce(h, labels)
                ce_losses.append(ce)
            all_rlogits.extend(rlogits)

        loss = None
        aux = None
        round_losses = None
        logits = None
        if labels is None:
            # 评测路径: 返回最终轮的 logits (训练路径 labels 恒非 None, 不受影响)
            logits = self.lm_head(self.model.norm(h)) / self.logits_scaling
        if labels is not None:
            loss = torch.stack(ce_losses).mean()
            round_losses = [c.detach() for c in ce_losses]  # 监控用: 各轮 CE 分开看
            if all_rlogits:
                aux = load_balancing_loss_func(
                    tuple(all_rlogits), self.num_experts, self.num_experts_per_tok)
                loss = loss + self.router_aux_loss_coef * aux
        return SimpleNamespace(loss=loss, logits=logits, aux_loss=aux, round_losses=round_losses)

    # ---------- 推理: KV cache 贪心生成 ----------
    # 思路: 循环架构下, 过去位置在第 r 轮第 l 层的注意力输入是静态的(因果性保证
    # 之后不再改变), 因此可以按 (round, layer) 缓存 K/V。预填充阶段复用标准前向,
    # 用 pre-hook 捕获每个注意力子层的输入并现算 K/V; 解码阶段每 token 只算当前
    # 位置, 数学上与整序列重算等价(仅浮点运算次序差异)。
    @torch.no_grad()
    def _prefill_kv(self, input_ids):
        """跑一遍标准前向, 返回 (最后一位的 logits, cache[r][l]=(K,V))。"""
        captured = []
        hooks = []
        for layer in self.model.layers:
            hooks.append(layer.self_attn.register_forward_pre_hook(
                lambda m, args, kwargs: captured.append(kwargs["hidden_states"].detach()),
                with_kwargs=True))
        try:
            out = self.forward(input_ids=input_ids)
        finally:
            for hk in hooks:
                hk.remove()
        L = len(self.model.layers)
        assert len(captured) == self.loop_rounds * L, \
            f"hook 捕获数 {len(captured)} != 轮数×层数 {self.loop_rounds * L}"
        from transformers.models.granitemoe.modeling_granitemoe import apply_rotary_pos_emb
        bsz, seq = input_ids.shape
        pos_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0)
        dummy = torch.zeros(1, seq, self.config.hidden_size, device=input_ids.device,
                            dtype=self.model.embed_tokens.weight.dtype)
        cos, sin = self.model.rotary_emb(dummy, pos_ids)
        cache = []
        for r in range(self.loop_rounds):
            rc = []
            for l, layer in enumerate(self.model.layers):
                x = captured[r * L + l]
                sa = layer.self_attn
                hd, kvh = sa.head_dim, self.config.num_key_value_heads
                k = sa.k_proj(x).view(bsz, seq, kvh, hd).transpose(1, 2)
                v = sa.v_proj(x).view(bsz, seq, kvh, hd).transpose(1, 2)
                _, k = apply_rotary_pos_emb(
                    torch.zeros(bsz, self.config.num_attention_heads, seq, hd,
                                device=x.device, dtype=k.dtype), k, cos, sin)
                rc.append((k, v))
            cache.append(rc)
        return out.logits[:, -1], cache

    def _attn_decode(self, layer, x, k_cache, v_cache, cos, sin):
        """解码阶段的单 token 注意力: x (b,1,d) 已加 cache 的 K/V (b,kvh,t,hd)。"""
        from transformers.models.granitemoe.modeling_granitemoe import apply_rotary_pos_emb
        sa = layer.self_attn
        b = x.shape[0]
        nh, kvh, hd = self.config.num_attention_heads, self.config.num_key_value_heads, sa.head_dim
        q = sa.q_proj(x).view(b, 1, nh, hd).transpose(1, 2)
        k = sa.k_proj(x).view(b, 1, kvh, hd).transpose(1, 2)
        v = sa.v_proj(x).view(b, 1, kvh, hd).transpose(1, 2)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        K = torch.cat([k_cache, k], dim=2)
        V = torch.cat([v_cache, v], dim=2)
        g = nh // kvh
        Kf = K.repeat_interleave(g, dim=1)
        Vf = V.repeat_interleave(g, dim=1)
        w = (q @ Kf.transpose(-1, -2)) * sa.scaling
        o = w.softmax(-1) @ Vf                       # (b,nh,1,hd)
        o = o.transpose(1, 2).reshape(b, 1, nh * hd)
        return sa.o_proj(o), K, V

    @torch.no_grad()
    def greedy_generate(self, input_ids, max_new_tokens=64, eos_token_id=None):
        """KV cache 贪心生成, 返回生成的 id 列表(不含 prompt)。仅支持 batch=1。"""
        assert input_ids.shape[0] == 1
        dev = input_ids.device
        layers = self.model.layers
        mu = self.config.residual_multiplier
        L = len(layers)

        logits, cache = self._prefill_kv(input_ids)
        nxt = logits[0].argmax().item()
        out_ids = []
        pos = input_ids.shape[1]

        for _ in range(max_new_tokens):
            if eos_token_id is not None and nxt == eos_token_id:
                break
            out_ids.append(nxt)
            pos_ids = torch.tensor([[pos]], device=dev)
            cos, sin = self.model.rotary_emb(
                torch.zeros(1, 1, self.config.hidden_size, device=dev,
                            dtype=self.model.embed_tokens.weight.dtype), pos_ids)
            h = self.model.embed_tokens(
                torch.tensor([[nxt]], device=dev)) * self.model.embedding_multiplier
            b0 = h
            prev_blocks = []
            for r in range(self.loop_rounds):
                if not self.use_attnres:
                    for l, layer in enumerate(layers):
                        f, K, V = self._attn_decode(
                            layer, layer.input_layernorm(h), cache[r][l][0], cache[r][l][1], cos, sin)
                        cache[r][l] = (K, V)
                        h = h + mu * f
                        f, _ = self._sublayer(layer, "moe", h, None, None)
                        h = h + mu * f
                else:
                    fixed = [b0] + (list(prev_blocks) if self.cross_round else [])
                    block_sums = []
                    for blk in range(self.attnres_blocks):
                        blk_layers = layers[blk * self.layers_per_block:(blk + 1) * self.layers_per_block]
                        q_blk = self.pseudo_queries[r,
                                                    blk * self.sub_per_block:(blk + 1) * self.sub_per_block]
                        cat = torch.stack(fixed)
                        inv_rms = torch.rsqrt(cat.float().pow(2).mean(-1) + 1e-6)
                        block_sum = None
                        for j, layer in enumerate(blk_layers):
                            for kind in ("attn", "moe"):
                                q = q_blk[2 * j] if kind == "attn" else q_blk[2 * j + 1]
                                block_sum, _ = self._attnres_step_decode(
                                    layer, kind, q, cat, inv_rms, block_sum, cos, sin, mu,
                                    cache, r, blk * self.layers_per_block + j)
                        fixed.append(block_sum)
                        block_sums.append(block_sum)
                    final_sources = [b0] + (list(prev_blocks) if self.cross_round else []) + block_sums
                    q_final = self.pseudo_queries[r, 2 * L]
                    h = self._aggregate(q_final, final_sources)
                if self.cross_round:
                    prev_blocks = prev_blocks + block_sums
            logits = self.lm_head(self.model.norm(h)) / self.logits_scaling
            nxt = logits[0, -1].argmax().item()
            pos += 1
        return out_ids

    def _attnres_step_decode(self, layer, kind, q, cat, inv_rms, block_sum, cos, sin, mu,
                             cache, r, l):
        """_attnres_step 的解码版: 聚合 + 单个子层(单 token, 注意力走 cache)。"""
        logits = torch.einsum("d,sbtd->sbt", q, cat) * inv_rms
        if block_sum is not None:
            bs_inv = torch.rsqrt(block_sum.float().pow(2).mean(-1) + 1e-6)
            bs_logit = torch.einsum("d,btd->bt", q, block_sum) * bs_inv
            alpha = torch.cat([logits, bs_logit.unsqueeze(0)], dim=0).softmax(dim=0)
        else:
            alpha = logits.softmax(dim=0)
        alpha = alpha.to(cat.dtype)
        if block_sum is not None:
            h = torch.einsum("sbt,sbtd->btd", alpha[:-1], cat) + alpha[-1].unsqueeze(-1) * block_sum
        else:
            h = torch.einsum("sbt,sbtd->btd", alpha, cat)
        if kind == "attn":
            f, K, V = self._attn_decode(
                layer, layer.input_layernorm(h), cache[r][l][0], cache[r][l][1], cos, sin)
            cache[r][l] = (K, V)
            rl = None
        else:
            f, rl = self._sublayer(layer, "moe", h, None, None)
        upd = mu * f
        return (upd if block_sum is None else block_sum + upd), rl