"""ECHO-1-nano model. Stage 1: a modern decoder-only transformer with
RMSNorm (pre-norm), RoPE, grouped-query attention with optional QK-norm, SwiGLU
feed-forward, weight-tied head, and a KV cache for generation.

Later stages add MLA, MoE, and MTP behind config toggles. Kept deliberately small
and explicit (manual attention rather than fused kernels) so each piece is legible.
"""

from __future__ import annotations

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

from config import NanoConfig


class RMSNorm(nn.Module):
    """Root-mean-square layer norm (no mean subtraction, no bias)."""

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight


def build_rope(seq: int, head_dim: int, theta: float, device, dtype):
    """Precompute RoPE cos/sin tables of shape (seq, head_dim)."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq, device=device).float()
    freqs = torch.outer(t, inv_freq)              # (seq, head_dim/2)
    emb = torch.cat([freqs, freqs], dim=-1)       # (seq, head_dim)
    return emb.cos().to(dtype), emb.sin().to(dtype)


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, offset: int = 0):
    """x: (B, H, T, hd). cos/sin: (S, hd). Applies rotation at positions
    [offset, offset+T)."""
    T = x.shape[-2]
    c = cos[offset:offset + T].unsqueeze(0).unsqueeze(0)   # (1,1,T,hd)
    s = sin[offset:offset + T].unsqueeze(0).unsqueeze(0)
    return x * c + _rotate_half(x) * s


class Attention(nn.Module):
    """Grouped-query attention. n_kv_heads <= n_heads share K/V across query-head
    groups. Optional per-head RMSNorm on Q and K (QK-norm). KV cache for decode."""

    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.n_kv = cfg.n_kv_heads
        self.hd = cfg.head_dim
        assert self.n_heads % self.n_kv == 0, "n_heads must be a multiple of n_kv_heads"
        self.rep = self.n_heads // self.n_kv

        self.wq = nn.Linear(cfg.d_model, self.n_heads * self.hd, bias=False)
        self.wk = nn.Linear(cfg.d_model, self.n_kv * self.hd, bias=False)
        self.wv = nn.Linear(cfg.d_model, self.n_kv * self.hd, bias=False)
        self.wo = nn.Linear(self.n_heads * self.hd, cfg.d_model, bias=False)
        self.qk_norm = cfg.qk_norm
        if cfg.qk_norm:
            self.q_norm = RMSNorm(self.hd)
            self.k_norm = RMSNorm(self.hd)
        self.drop = cfg.dropout

    def forward(self, x, cos, sin, cache=None):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)  # (B,H,T,hd)
        k = self.wk(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)     # (B,Hkv,T,hd)
        v = self.wv(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)

        if self.qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)

        offset = 0 if cache is None else cache[0].shape[-2]
        q = apply_rope(q, cos, sin, offset)
        k = apply_rope(k, cos, sin, offset)

        if cache is not None:
            pk, pv = cache
            k = torch.cat([pk, k], dim=-2)
            v = torch.cat([pv, v], dim=-2)
        new_cache = (k, v)

        # expand KV groups to full head count
        k = k.repeat_interleave(self.rep, dim=1)
        v = v.repeat_interleave(self.rep, dim=1)

        # causal only needed when not decoding a single step against a full cache
        is_causal = cache is None or T > 1
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=is_causal,
            dropout_p=self.drop if self.training else 0.0,
        )
        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.hd)
        return self.wo(out), new_cache


class MLAttention(nn.Module):
    """Multi-head latent attention (DeepSeek-V2). K/V are compressed into a shared
    low-rank latent c_kv that is the only thing cached, then up-projected to a
    per-head 'nope' (no-position) part. Position is carried by a separate decoupled
    RoPE key k_rope shared across heads. Query is split the same way and may also be
    low-rank compressed. The KV cache shrinks from n_kv*head_dim to
    kv_lora_rank + rope_head_dim per token."""

    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.nope = cfg.head_dim                                  # per-head content dim
        self.rope = cfg.rope_head_dim or (cfg.head_dim // 2)      # decoupled position dim
        self.vhd = cfg.head_dim
        self.qdim = self.nope + self.rope
        self.scale = self.qdim ** -0.5
        self.kv_rank = cfg.kv_lora_rank
        self.q_rank = cfg.q_lora_rank

        # KV down/up
        self.w_dkv = nn.Linear(cfg.d_model, self.kv_rank, bias=False)     # compress
        self.kv_norm = RMSNorm(self.kv_rank)
        self.w_uk = nn.Linear(self.kv_rank, self.n_heads * self.nope, bias=False)
        self.w_uv = nn.Linear(self.kv_rank, self.n_heads * self.vhd, bias=False)
        self.w_kr = nn.Linear(cfg.d_model, self.rope, bias=False)         # shared rope key
        # Q (optionally compressed)
        if self.q_rank:
            self.w_dq = nn.Linear(cfg.d_model, self.q_rank, bias=False)
            self.q_norm = RMSNorm(self.q_rank)
            self.w_uq = nn.Linear(self.q_rank, self.n_heads * self.qdim, bias=False)
        else:
            self.w_q = nn.Linear(cfg.d_model, self.n_heads * self.qdim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.vhd, cfg.d_model, bias=False)
        self.drop = cfg.dropout

    def forward(self, x, cos, sin, cache=None):
        B, T, _ = x.shape
        H, nope, rope, vhd = self.n_heads, self.nope, self.rope, self.vhd

        # queries
        if self.q_rank:
            q = self.w_uq(self.q_norm(self.w_dq(x)))
        else:
            q = self.w_q(x)
        q = q.view(B, T, H, self.qdim).transpose(1, 2)           # (B,H,T,qdim)
        q_nope, q_rope = q[..., :nope], q[..., nope:]

        # compressed latent + decoupled rope key (these get cached)
        c_kv = self.kv_norm(self.w_dkv(x))                       # (B,T,rank)
        k_rope = self.w_kr(x).view(B, T, 1, rope).transpose(1, 2)  # (B,1,T,rope)

        offset = 0 if cache is None else cache[0].shape[1]
        q_rope = apply_rope(q_rope, cos, sin, offset)
        k_rope = apply_rope(k_rope, cos, sin, offset)

        if cache is not None:
            pc, pkr = cache                                      # (B,Tp,rank),(B,1,Tp,rope)
            c_kv = torch.cat([pc, c_kv], dim=1)
            k_rope = torch.cat([pkr, k_rope], dim=2)
        new_cache = (c_kv, k_rope)
        S = c_kv.shape[1]

        # up-project the (full) latent to per-head K-nope and V
        k_nope = self.w_uk(c_kv).view(B, S, H, nope).transpose(1, 2)   # (B,H,S,nope)
        v = self.w_uv(c_kv).view(B, S, H, vhd).transpose(1, 2)         # (B,H,S,vhd)
        k_rope = k_rope.expand(B, H, S, rope)                          # share across heads

        q_full = torch.cat([q_nope, q_rope], dim=-1)            # (B,H,T,qdim)
        k_full = torch.cat([k_nope, k_rope], dim=-1)            # (B,H,S,qdim)

        is_causal = cache is None or T > 1
        out = F.scaled_dot_product_attention(
            q_full, k_full, v, is_causal=is_causal,
            dropout_p=self.drop if self.training else 0.0,
        )
        out = out.transpose(1, 2).contiguous().view(B, T, H * vhd)
        return self.wo(out), new_cache


class SwiGLU(nn.Module):
    """Gated FFN: (silu(W1 x) * W3 x) W2."""

    def __init__(self, d: int, h: int):
        super().__init__()
        self.w1 = nn.Linear(d, h, bias=False)
        self.w3 = nn.Linear(d, h, bias=False)
        self.w2 = nn.Linear(h, d, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class DenseFFN(nn.Module):
    """A single SwiGLU FFN. Returns (out, aux=0) to match the MoE interface."""

    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.ff = SwiGLU(cfg.d_model, cfg.ff_dim)

    def forward(self, x):
        return self.ff(x), x.new_zeros(())


class MoE(nn.Module):
    """Sparse mixture of experts (DeepSeek-V3 flavour): top-k routed fine-grained
    experts + always-on shared experts, sigmoid gating, and aux-loss-free load
    balancing via a per-expert selection bias that drifts toward equal load. An
    optional classic load-balance aux loss is available via moe_aux_alpha."""

    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.n_exp = cfg.n_experts
        self.top_k = cfg.top_k
        self.alpha = cfg.moe_aux_alpha
        self.experts = nn.ModuleList(SwiGLU(cfg.d_model, cfg.ff_dim) for _ in range(cfg.n_experts))
        self.shared = nn.ModuleList(SwiGLU(cfg.d_model, cfg.ff_dim) for _ in range(cfg.n_shared))
        self.router = nn.Linear(cfg.d_model, cfg.n_experts, bias=False)
        # balancing bias: affects selection only, not the gate weights. Updated as a
        # buffer (no gradient), DeepSeek-V3 style.
        self.register_buffer("bias", torch.zeros(cfg.n_experts))
        self.bias_lr = 1e-3

    def forward(self, x):
        B, T, d = x.shape
        xf = x.reshape(-1, d)                                  # (N,d)
        N = xf.shape[0]
        gates = torch.sigmoid(self.router(xf))                # (N,E)
        # select top-k by gate + balancing bias (bias steers load, not weights)
        sel = torch.topk(gates + self.bias, self.top_k, dim=-1).indices   # (N,k)
        sel_gate = torch.gather(gates, 1, sel)                 # (N,k)
        denom = sel_gate.sum(-1, keepdim=True).clamp_min(1e-9)
        weight = sel_gate / denom                             # normalized over chosen

        out = torch.zeros_like(xf)
        load = torch.zeros(self.n_exp, device=xf.device)
        for e in range(self.n_exp):
            hit = sel == e                                     # (N,k) bool
            tok = hit.any(-1)                                  # (N,)
            load[e] = tok.sum()
            if tok.any():
                we = (weight * hit).sum(-1)[tok].unsqueeze(-1)  # (n_e,1)
                out[tok] += we * self.experts[e](xf[tok])
        for se in self.shared:
            out = out + se(xf)

        aux = xf.new_zeros(())
        if self.training:
            # aux-loss-free: nudge bias so underused experts become more selectable
            with torch.no_grad():
                target = N * self.top_k / self.n_exp
                self.bias += self.bias_lr * (target - load).sign()
            if self.alpha > 0:
                f = load / (N * self.top_k)                    # fraction of slots
                p = gates.mean(0)                              # mean gate per expert
                aux = self.alpha * self.n_exp * (f * p).sum()
        return out.view(B, T, d), aux


class Block(nn.Module):
    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model)
        self.attn = MLAttention(cfg) if cfg.attn == "mla" else Attention(cfg)
        self.ffn_norm = RMSNorm(cfg.d_model)
        self.ffn = MoE(cfg) if cfg.ffn == "moe" else DenseFFN(cfg)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, cos, sin, cache=None):
        h, new_cache = self.attn(self.attn_norm(x), cos, sin, cache)
        x = x + self.drop(h)
        f, aux = self.ffn(self.ffn_norm(x))
        x = x + self.drop(f)
        return x, new_cache, aux


class MTPModule(nn.Module):
    """One multi-token-prediction depth (DeepSeek-V3). Combines the previous
    depth's hidden state with the embedding of the next observed token, then runs
    a transformer block. The shared output head turns its hidden into logits for a
    token one step further ahead."""

    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.h_norm = RMSNorm(cfg.d_model)
        self.e_norm = RMSNorm(cfg.d_model)
        self.proj = nn.Linear(2 * cfg.d_model, cfg.d_model, bias=False)
        self.block = Block(cfg)

    def forward(self, h_prev, emb_next, cos, sin):
        x = self.proj(torch.cat([self.h_norm(h_prev), self.e_norm(emb_next)], dim=-1))
        x, _, _ = self.block(x, cos, sin, None)
        return x


def _shift_left(t: torch.Tensor, k: int, fill):
    """Move positions left by k: out[:, i] = t[:, i+k]; last k filled with `fill`."""
    B, T = t.shape[0], t.shape[1]
    tail_shape = (B, k) + tuple(t.shape[2:])
    pad = t.new_full(tail_shape, fill) if t.dim() == 2 else t.new_zeros(tail_shape)
    return torch.cat([t[:, k:], pad], dim=1)


class NanoLM(nn.Module):
    def __init__(self, cfg: NanoConfig):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.head.weight = self.embed.weight
        self.mtp = nn.ModuleList(MTPModule(cfg) for _ in range(cfg.mtp_tokens))
        # RoPE dim: GQA rotates the whole head; MLA rotates only the decoupled part.
        self.rope_dim = cfg.head_dim if cfg.attn == "gqa" else (cfg.rope_head_dim or cfg.head_dim // 2)
        self._cos = self._sin = None
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def _rope(self, device, dtype):
        if self._cos is None or self._cos.device != device or self._cos.dtype != dtype:
            self._cos, self._sin = build_rope(self.cfg.max_seq, self.rope_dim,
                                              self.cfg.rope_theta, device, dtype)
        return self._cos, self._sin

    def forward(self, ids, targets=None, caches=None):
        emb = self.embed(ids)
        x = emb
        cos, sin = self._rope(x.device, x.dtype)
        new_caches = []
        aux_total = x.new_zeros(())
        for i, blk in enumerate(self.blocks):
            c = None if caches is None else caches[i]
            x, nc, aux = blk(x, cos, sin, c)
            new_caches.append(nc)
            aux_total = aux_total + aux
        trunk = x                                   # pre-final-norm hidden (MTP root)
        logits = self.head(self.norm(trunk))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1),
                                   ignore_index=-100) + aux_total
            # multi-token prediction: depth k predicts the token k steps further ahead
            if self.mtp:
                h_prev = trunk
                mtp_loss = x.new_zeros(())
                for k in range(1, len(self.mtp) + 1):
                    emb_next = self.embed(_shift_left(ids, k, 0))      # token at i+k
                    h_prev = self.mtp[k - 1](h_prev, emb_next, cos, sin)
                    lg = self.head(self.norm(h_prev))
                    tgt = _shift_left(targets, k, -100)                # token at i+1+k
                    mtp_loss = mtp_loss + F.cross_entropy(
                        lg.reshape(-1, lg.size(-1)), tgt.reshape(-1), ignore_index=-100)
                loss = loss + self.cfg.mtp_lambda / len(self.mtp) * mtp_loss
        return logits, loss, new_caches

    @torch.no_grad()
    def generate(self, ids, max_new: int, temperature: float = 1.0, top_k: int = 0,
                 eos_id: int | None = None):
        self.eval()
        caches = None
        # prime the cache with the prompt
        logits, _, caches = self.forward(ids, caches=None)
        out = ids
        for _ in range(max_new):
            last = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k:
                v, _ = torch.topk(last, top_k)
                last = last.masked_fill(last < v[:, [-1]], float("-inf"))
            probs = F.softmax(last, dim=-1)
            nxt = torch.multinomial(probs, 1)
            out = torch.cat([out, nxt], dim=1)
            if eos_id is not None and out.size(0) == 1 and int(nxt.item()) == eos_id:
                break
            logits, _, caches = self.forward(nxt, caches=caches)
        return out

    def n_params(self) -> int:
        n = sum(p.numel() for p in self.parameters())
        if self.cfg.tie_embeddings:
            n -= self.embed.weight.numel()  # counted once
        return n