"""VBQ transformer: Llama-style decoder (RMSNorm, RoPE, GQA, QK-norm, SwiGLU, tied embeddings)."""
from __future__ import annotations

import math
from dataclasses import dataclass, asdict

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


@dataclass
class VBQConfig:
    vocab_size: int = 32768
    d_model: int = 640
    n_layers: int = 12
    n_heads: int = 10
    n_kv_heads: int = 5
    mlp_hidden: int = 1664
    max_seq_len: int = 2048
    rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    qk_norm: bool = True
    tie_embeddings: bool = True
    # Mixture of experts (0 = dense SwiGLU). Top-1 routing, Switch-style.
    n_experts: int = 0
    capacity_factor: float = 1.25

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads

    def n_params(self) -> int:
        d, hd = self.d_model, self.head_dim
        attn = d * d * 2 + d * (self.n_kv_heads * hd) * 2
        mlp = 3 * d * self.mlp_hidden * max(1, self.n_experts) + (d * self.n_experts if self.n_experts else 0)
        per_layer = attn + mlp + 2 * d + (2 * hd if self.qk_norm else 0)
        emb = self.vocab_size * d * (1 if self.tie_embeddings else 2)
        return emb + self.n_layers * per_layer + d

    def to_dict(self):
        return asdict(self)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps)


def build_rope(seq_len: int, head_dim: int, theta: float, device=None):
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    t = torch.arange(seq_len, device=device, dtype=torch.float32)
    freqs = torch.outer(t, inv)
    return freqs.cos(), freqs.sin()  # [T, hd/2]


def apply_rope(x, cos, sin):
    # x: [B, H, T, hd]; half-split rotation
    x1, x2 = x.chunk(2, dim=-1)
    cos = cos[None, None, : x.shape[2]].to(x.dtype)
    sin = sin[None, None, : x.shape[2]].to(x.dtype)
    return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)


class Attention(nn.Module):
    def __init__(self, c: VBQConfig):
        super().__init__()
        self.c = c
        hd = c.head_dim
        self.wq = nn.Linear(c.d_model, c.n_heads * hd, bias=False)
        self.wk = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False)
        self.wv = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False)
        self.wo = nn.Linear(c.n_heads * hd, c.d_model, bias=False)
        if c.qk_norm:
            self.q_norm = RMSNorm(hd, c.norm_eps)
            self.k_norm = RMSNorm(hd, c.norm_eps)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        c, hd = self.c, self.c.head_dim
        q = self.wq(x).view(B, T, c.n_heads, hd).transpose(1, 2)
        k = self.wk(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2)
        v = self.wv(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2)
        if c.qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if c.n_kv_heads != c.n_heads:
            rep = c.n_heads // c.n_kv_heads
            k = k[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd)
            v = v[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.wo(y.transpose(1, 2).reshape(B, T, c.n_heads * hd))


class MLP(nn.Module):
    def __init__(self, c: VBQConfig):
        super().__init__()
        self.gate_up = nn.Linear(c.d_model, 2 * c.mlp_hidden, bias=False)
        self.down = nn.Linear(c.mlp_hidden, c.d_model, bias=False)

    def forward(self, x):
        g, u = self.gate_up(x).chunk(2, dim=-1)
        return self.down(F.silu(g) * u)


class MoE(nn.Module):
    """Top-1 mixture of SwiGLU experts (Switch Transformer routing).

    Fixed shapes (compile-friendly): every expert gets a buffer of
    capacity = ceil(capacity_factor * tokens / experts) slots while training
    (tokens past it skip the MLP and keep only the residual); at eval every
    token fits. The chosen expert's output is scaled by p / p.detach(), which
    is 1 in value, so a model upcycled from identical copies computes exactly
    what the dense model did, while the router still gets gradients.
    Returns (y, aux) with aux = Switch load-balancing loss (1.0 = balanced).
    """

    def __init__(self, c: VBQConfig):
        super().__init__()
        self.c = c
        E, d, h = c.n_experts, c.d_model, c.mlp_hidden
        self.router = nn.Linear(d, E, bias=False)
        self.gate_up = nn.Parameter(torch.empty(E, 2 * h, d))
        self.down = nn.Parameter(torch.empty(E, d, h))
        nn.init.normal_(self.gate_up, 0.0, 0.02)
        nn.init.normal_(self.down, 0.0, 0.02)

    def forward(self, x):
        shape = x.shape
        x = x.reshape(-1, shape[-1])
        N, d = x.shape
        E = self.c.n_experts
        probs = F.softmax(self.router(x).float(), dim=-1)          # [N, E]
        gate, idx = probs.max(-1)                                  # [N]
        onehot = F.one_hot(idx, E)                                 # [N, E]
        pos = ((onehot.cumsum(0) - 1) * onehot).sum(-1)            # slot within its expert
        C = math.ceil(self.c.capacity_factor * N / E) if self.training else N
        keep = pos < C
        slot = torch.where(keep, idx * C + pos, torch.full_like(pos, E * C))  # overflow -> dump slot
        buf = x.new_zeros(E * C + 1, d).index_put((slot,), x)
        hbuf = buf[: E * C].view(E, C, d)
        g, u = torch.bmm(hbuf, self.gate_up.to(x.dtype).transpose(1, 2)).chunk(2, dim=-1)
        y = torch.bmm(F.silu(g) * u, self.down.to(x.dtype).transpose(1, 2))   # [E, C, d]
        y = torch.cat([y.reshape(E * C, d), y.new_zeros(1, d)])[slot]
        scale = (gate / gate.detach()) * keep
        y = y * scale.to(y.dtype).unsqueeze(-1)
        aux = E * (onehot.float().mean(0) * probs.mean(0)).sum()
        return y.reshape(shape), aux


class Block(nn.Module):
    def __init__(self, c: VBQConfig):
        super().__init__()
        self.attn_norm = RMSNorm(c.d_model, c.norm_eps)
        self.attn = Attention(c)
        self.mlp_norm = RMSNorm(c.d_model, c.norm_eps)
        self.mlp = MoE(c) if c.n_experts else MLP(c)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        return x + self.ffn(self.mlp_norm(x))[0]

    def ffn(self, x):
        """(output, load-balancing loss) for dense and MoE blocks alike."""
        y = self.mlp(x)
        return y if isinstance(y, tuple) else (y, x.new_zeros((), dtype=torch.float32))


class FusedLinearCE(torch.autograd.Function):
    """Cross-entropy over a huge vocab without ever materialising full logits.

    Gradients w.r.t. hidden states and the (tied) output matrix are computed chunk by chunk
    during the forward pass, so the backward pass is a cheap scale. Saves ~4 GB at 32k vocab.
    """

    @staticmethod
    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
        N = h.shape[0]
        wb = w.to(h.dtype)
        grad_h = torch.empty_like(h)
        grad_w = torch.zeros_like(w, dtype=torch.float32)
        total = torch.zeros((), device=h.device, dtype=torch.float32)
        for i in range(0, N, chunk):
            hc, tc = h[i : i + chunk], targets[i : i + chunk]
            logits = (hc @ wb.t()).float()
            lse = torch.logsumexp(logits, dim=-1)
            tgt = logits.gather(1, tc[:, None]).squeeze(1)
            total += (lse - tgt).sum()
            p = torch.exp(logits - lse[:, None])
            if z_loss:
                total += z_loss * (lse * lse).sum()
                p.mul_(1.0 + 2.0 * z_loss * lse[:, None])
            p.scatter_add_(1, tc[:, None], torch.full_like(tc[:, None], -1, dtype=p.dtype))
            p.mul_(1.0 / N)
            pb = p.to(h.dtype)
            grad_h[i : i + chunk] = pb @ wb
            grad_w += (pb.t() @ hc).float()
        ctx.save_for_backward(grad_h, grad_w)
        return total / N

    @staticmethod
    def backward(ctx, g):
        grad_h, grad_w = ctx.saved_tensors
        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None


class VBQ(nn.Module):
    def __init__(self, c: VBQConfig):
        super().__init__()
        self.c = c
        self.tok_emb = nn.Embedding(c.vocab_size, c.d_model)
        self.blocks = nn.ModuleList(Block(c) for _ in range(c.n_layers))
        self.final_norm = RMSNorm(c.d_model, c.norm_eps)
        self.lm_head = None if c.tie_embeddings else nn.Linear(c.d_model, c.vocab_size, bias=False)
        cos, sin = build_rope(c.max_seq_len, c.head_dim, c.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(self._init)
        for n, p in self.named_parameters():  # scaled init on residual output projections
            if n.endswith("wo.weight") or n.endswith("down.weight"):
                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * c.n_layers))

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)

    @property
    def out_weight(self):
        return self.tok_emb.weight if self.lm_head is None else self.lm_head.weight

    def hidden(self, idx):
        return self.hidden_aux(idx)[0]

    def hidden_aux(self, idx):
        """Final hidden states and the mean MoE load-balancing loss (0 for dense)."""
        return self._blocks(self.tok_emb(idx))

    def hidden_patch(self, idx, k: int):
        """Patch-level training (Shao et al. 2024): idx [B, P*k]; each group of k tokens enters as the mean of its
        token embeddings, so the transformer runs on P positions. Same weights as token-level training."""
        B, L = idx.shape
        return self._blocks(self.tok_emb(idx).view(B, L // k, k, -1).mean(2))

    def _blocks(self, x):
        aux = x.new_zeros((), dtype=torch.float32)
        for blk in self.blocks:
            x = x + blk.attn(blk.attn_norm(x), self.rope_cos, self.rope_sin)
            y, a = blk.ffn(blk.mlp_norm(x))
            x = x + y
            aux = aux + a
        return self.final_norm(x), aux / len(self.blocks)

    def forward(self, idx, targets=None, z_loss: float = 0.0, ce_chunk: int = 4096):
        h = self.hidden(idx)
        if targets is None:
            return (h @ self.out_weight.to(h.dtype).t()).float()
        loss = FusedLinearCE.apply(h.reshape(-1, h.shape[-1]), self.out_weight, targets.reshape(-1), ce_chunk, z_loss)
        return loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0):
        self.eval()
        for _ in range(max_new_tokens):
            ctx = idx[:, -self.c.max_seq_len :]
            logits = self(ctx)[:, -1]
            if repetition_penalty != 1.0:
                for tok in set(ctx[0].tolist()):
                    logits[0, tok] = logits[0, tok] / repetition_penalty if logits[0, tok] > 0 else logits[0, tok] * repetition_penalty
            if temperature <= 0:
                nxt = logits.argmax(-1, keepdim=True)
            else:
                logits = logits / temperature
                if top_k:
                    v, _ = torch.topk(logits, min(top_k, logits.shape[-1]))
                    logits[logits < v[:, [-1]]] = -float("inf")
                probs = F.softmax(logits, dim=-1)
                if top_p < 1.0:
                    sp, si = torch.sort(probs, descending=True)
                    cum = sp.cumsum(-1)
                    sp[cum - sp > top_p] = 0
                    probs = torch.zeros_like(probs).scatter_(1, si, sp)
                    probs = probs / probs.sum(-1, keepdim=True)
                nxt = torch.multinomial(probs, 1)
            idx = torch.cat([idx, nxt], dim=1)
            yield int(nxt.item())
            if eos_id is not None and int(nxt.item()) == eos_id:
                break

    @torch.no_grad()
    def generate_cached(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0):
        """KV-cached sampling. Yields token ids one by one. idx: [1, T] prompt."""
        self.eval()
        c = self.c
        idx = idx[:, -(c.max_seq_len - 1):]
        stop_ids = set() if eos_id is None else ({eos_id} if isinstance(eos_id, int) else set(eos_id))
        caches: list = [None] * c.n_layers
        seen = set(idx[0].tolist())

        def forward(ids, pos):
            T = ids.shape[1]
            x = self.tok_emb(ids)
            cos, sin = self.rope_cos[pos : pos + T], self.rope_sin[pos : pos + T]
            for i, blk in enumerate(self.blocks):
                a = blk.attn
                h = blk.attn_norm(x)
                q = a.wq(h).view(1, T, c.n_heads, c.head_dim).transpose(1, 2)
                k = a.wk(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2)
                v = a.wv(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2)
                if c.qk_norm:
                    q, k = a.q_norm(q), a.k_norm(k)
                q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
                if caches[i] is not None:
                    k, v = torch.cat([caches[i][0], k], 2), torch.cat([caches[i][1], v], 2)
                caches[i] = (k, v)
                rep = c.n_heads // c.n_kv_heads
                kk = k.repeat_interleave(rep, dim=1) if rep > 1 else k
                vv = v.repeat_interleave(rep, dim=1) if rep > 1 else v
                y = F.scaled_dot_product_attention(q, kk, vv, is_causal=(T > 1))
                x = x + a.wo(y.transpose(1, 2).reshape(1, T, c.n_heads * c.head_dim))
                x = x + blk.ffn(blk.mlp_norm(x))[0]
            return (self.final_norm(x[:, -1]) @ self.out_weight.to(x.dtype).t()).float()

        logits = forward(idx, 0)
        pos = idx.shape[1]
        for _ in range(max_new_tokens):
            if pos >= c.max_seq_len:
                break
            lg = logits[0].clone()
            if repetition_penalty != 1.0:
                for tok in seen:
                    lg[tok] = lg[tok] / repetition_penalty if lg[tok] > 0 else lg[tok] * repetition_penalty
            if temperature <= 0:
                nxt = int(lg.argmax())
            else:
                lg = lg / temperature
                if top_k:
                    v, _ = torch.topk(lg, min(top_k, lg.numel()))
                    lg[lg < v[-1]] = -float("inf")
                probs = F.softmax(lg, dim=-1)
                if top_p < 1.0:
                    sp, si = torch.sort(probs, descending=True)
                    sp[sp.cumsum(-1) - sp > top_p] = 0
                    probs = torch.zeros_like(probs).scatter_(0, si, sp)
                    probs = probs / probs.sum()
                nxt = int(torch.multinomial(probs, 1))
            yield nxt
            if nxt in stop_ids:
                break
            seen.add(nxt)
            logits = forward(torch.tensor([[nxt]], device=idx.device), pos)
            pos += 1