"""
Morphable Hybrid Attention — layer-level linearization with learnable selection,
combined with TriAttention token-level KV compression on the surviving full layers.

Implements the four-stage pipeline from the "Morphable Layer" figure:

  (1) Morphable Layer Construction
      Every attention layer gets a frozen Full-Attention path A_full (the pretrained
      MLADerfXSAAttention) and a trainable Linear-Attention sibling A_lin that reuses
      the frozen Q/K/V/O projections but replaces softmax with a learnable positive
      feature map (LoLCATs-style). A_lin is trained to match A_full per layer:
          L_hidden = (1/L) Σ_l || H_lin^(l) - H_full^(l) ||²

  (2) Layer Selection via Joint Opt. + Linearization Reg.
      Each layer mixes the two paths through a learnable gate α^(l) ∈ (0,1):
          H_mix^(l) = α^(l) · H_full^(l) + (1 - α^(l)) · H_lin^(l)
      trained on synthetic passkey-retrieval data with
          L_total = L_align + λ L_reg
                  = (1/(L|T|)) Σ_l Σ_t || H_mix,t^(l) - H_full,t^(l) ||²  +  λ Σ_l α^(l)
      The α-penalty pushes layers toward linear unless full attention is genuinely
      needed (retrieval heads), so α ranks layer importance.

  (3) Discretize Hybrid Layers
      Keep the top-k layers by α as Full Attention; linearize the rest.

  (4) Distillation & Finetuning
      Logits KL distillation (student hybrid vs. frozen teacher) + long-context FT.

Combined with triattention.py: the discretized "full" layers still hold a KV cache,
so TriAttention prunes them to a token budget; the "linear" layers carry an O(1)
recurrent state instead. See model.py generate_hybrid().
"""

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


# --------------------------------------------------------------------------- #
# trainable linear attention (reuses the frozen full-attention projections)
# --------------------------------------------------------------------------- #
def _feature_map(x, temp):
    """Positive feature map φ(x) = elu(x·temp) + 1 (LoLCATs-style, learnable temp)."""
    return F.elu(x * temp) + 1.0


class LinearAttention(nn.Module):
    """Linear-attention twin of MLADerfXSAAttention.

    Shares (by reference) the frozen full-attention Q/K/V/O projections, QK-norm and
    RoPE, and only learns a small per-head feature-map temperature. Softmax(qkᵀ) is
    replaced by the kernel weight φ(q)·φ(k), giving causal linear attention with an
    O(d²) recurrent state for decoding:

        out_t = Σ_{s≤t} (φ(q_t)·φ(k_s)) v_s  /  Σ_{s≤t} (φ(q_t)·φ(k_s))
    """

    def __init__(self, full_attn, cfg):
        super().__init__()
        self.full = full_attn                       # frozen; used only for its weights
        self.num_heads = cfg.n_head
        self.num_kv_heads = cfg.num_key_value_heads
        self.head_dim = cfg.head_dim
        self.nope_head_dim = cfg.nope_head_dim
        self.kv_groups = self.num_heads // self.num_kv_heads
        self.use_qk_norm = cfg.use_qk_norm
        # learnable per-head feature-map temperatures (init 1 -> φ = elu+1)
        self.q_temp = nn.Parameter(torch.ones(self.num_heads, 1))
        self.k_temp = nn.Parameter(torch.ones(self.num_kv_heads, 1))

    def _project(self, x, position_ids):
        """Reuse the frozen full path's projections + RoPE to get post-RoPE q,k,v."""
        f = self.full
        B, S, _ = x.shape
        q = f.q_b_proj(f.q_a_norm(f.q_a_proj(x)))
        q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        k = f.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = f.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
        if self.use_qk_norm:
            q, k = f.q_norm(q), f.k_norm(k)
        d = self.nope_head_dim
        q = torch.cat([q[..., :d], f.rope(q[..., d:], position_ids)], dim=-1)
        k = torch.cat([k[..., :d], f.rope(k[..., d:], position_ids)], dim=-1)
        return q, k, v

    def _out(self, y):
        f = self.full
        B, H, S, D = y.shape
        y = y.transpose(1, 2).contiguous().view(B, S, H * D)
        return f.o_b_proj(f.o_a_proj(y))

    def forward(self, x, position_ids, past_kv=None, use_cache=False):
        f = self.full
        B, S, _ = x.shape
        q, k, v = self._project(x, position_ids)
        # expand kv heads to query heads (GQA)
        if self.kv_groups > 1:
            k = k.repeat_interleave(self.kv_groups, dim=1)
            v = v.repeat_interleave(self.kv_groups, dim=1)
        qt = self.q_temp.repeat_interleave(1, 0).view(1, self.num_heads, 1, 1)
        kt = self.k_temp.repeat_interleave(self.kv_groups, 0).view(1, self.num_heads, 1, 1)
        phi_q = _feature_map(q, qt)                 # [B,H,S,D]
        phi_k = _feature_map(k, kt)

        # recurrent decoding: carry state (KV = Σ φk⊗v, Z = Σ φk)
        if use_cache and past_kv is not None:
            state_kv, state_z = past_kv[1], past_kv[2]   # ('linear', KV, Z)
            # accumulate this step's tokens into the state, then read out causally
            outs = []
            for t in range(S):
                pk, vv = phi_k[:, :, t], v[:, :, t]      # [B,H,D]
                state_kv = state_kv + pk.unsqueeze(-1) * vv.unsqueeze(-2)  # [B,H,D,D]
                state_z = state_z + pk                                     # [B,H,D]
                pq = phi_q[:, :, t]
                num = (pq.unsqueeze(-2) @ state_kv).squeeze(-2)            # [B,H,D]
                den = (pq * state_z).sum(-1, keepdim=True).clamp_min(1e-6)
                outs.append(num / den)
            y = torch.stack(outs, dim=2)                 # [B,H,S,D]
            present = ("linear", state_kv, state_z)
            return self._out(y), present

        # parallel form (exact same result), O(S²) — used for training / prefill
        w = torch.matmul(phi_q, phi_k.transpose(-2, -1))     # [B,H,S,S] kernel weights
        offset = 0 if past_kv is None else 0
        qpos = torch.arange(S, device=x.device).view(S, 1)
        kpos = torch.arange(S, device=x.device).view(1, S)
        w = w.masked_fill((kpos > qpos).unsqueeze(0).unsqueeze(0), 0.0)
        w = w / w.sum(-1, keepdim=True).clamp_min(1e-6)
        y = torch.matmul(w, v)                               # [B,H,S,D]
        if use_cache:
            # build the recurrent state from the full prefix for later decoding
            state_kv = torch.einsum("bhsd,bhse->bhde", phi_k, v)   # [B,H,D,D]
            state_z = phi_k.sum(2)                                 # [B,H,D]
            return self._out(y), ("linear", state_kv, state_z)
        return self._out(y)


# --------------------------------------------------------------------------- #
# morphable wrapper: full + linear + learnable gate α
# --------------------------------------------------------------------------- #
class MorphableAttention(nn.Module):
    """Wraps the pretrained full attention with a linear twin and a per-layer gate.

    mode:
        'mix'    — H = α·H_full + (1-α)·H_lin   (stages 1-2; captures alignment loss)
        'full'   — H = H_full                    (discretized: selected layer)
        'linear' — H = H_lin                     (discretized: linearized layer)
    """

    def __init__(self, full_attn, cfg, alpha_init=0.5):
        super().__init__()
        self.full = full_attn
        self.lin = LinearAttention(full_attn, cfg)
        # gate stored as a logit; α = sigmoid(logit)
        self.alpha_logit = nn.Parameter(torch.tensor(math.log(alpha_init / (1 - alpha_init))))
        self.mode = "mix"
        self.last_hidden_align = None     # ||H_lin - H_full||² captured on the last forward

    @property
    def alpha(self):
        return torch.sigmoid(self.alpha_logit)

    def freeze_full(self):
        for p in self.full.parameters():
            p.requires_grad_(False)

    def forward(self, x, position_ids, past_kv=None, use_cache=False):
        if self.mode == "full":
            return self.full(x, position_ids, past_kv=past_kv, use_cache=use_cache)
        if self.mode == "linear":
            return self.lin(x, position_ids, past_kv=past_kv, use_cache=use_cache)

        # 'mix': run both paths (no cache during training/selection)
        h_full = self.full(x, position_ids)
        h_lin = self.lin(x, position_ids)
        self.last_hidden_align = ((h_lin - h_full) ** 2).mean()
        a = self.alpha
        h_mix = a * h_full + (1.0 - a) * h_lin
        if use_cache:
            return h_mix, None
        return h_mix


# --------------------------------------------------------------------------- #
# stage-3 discretization + loss helpers
# --------------------------------------------------------------------------- #
@torch.no_grad()
def discretize(model, k_full):
    """Keep the top-k layers by α as full attention; linearize the rest.
    Returns the list of selected (full) layer indices."""
    morphs = [layer.attn for layer in model.layers
              if isinstance(layer.attn, MorphableAttention)]
    alphas = torch.stack([m.alpha.detach() for m in morphs])
    keep = set(torch.topk(alphas, min(k_full, len(morphs))).indices.tolist())
    for i, m in enumerate(morphs):
        m.mode = "full" if i in keep else "linear"
    return sorted(keep)


def hidden_alignment_loss(model):
    """L_hidden = mean_l ||H_lin^(l) - H_full^(l)||²  (stage 1)."""
    terms = [layer.attn.last_hidden_align for layer in model.layers
             if isinstance(layer.attn, MorphableAttention)
             and layer.attn.last_hidden_align is not None]
    if not terms:
        return None
    return torch.stack(terms).mean()


def linearization_reg(model):
    """L_reg = Σ_l α^(l)  (stage 2 penalty that pushes layers toward linear)."""
    terms = [layer.attn.alpha for layer in model.layers
             if isinstance(layer.attn, MorphableAttention)]
    if not terms:
        return None
    return torch.stack(terms).sum()


def set_mode(model, mode):
    for layer in model.layers:
        if isinstance(layer.attn, MorphableAttention):
            layer.attn.mode = mode