"""TinyLiquid -- our own tiny liquid-architecture language model.

Non-transformer design (no attention):
  * liquid blocks, each = basis expansion layer + gated MLP (dense or MoE),
    both with a sigmoid forget gate, residual connections, RMSNorm.
  * basis expansion: expand d -> N*B, group-norm within each basis block,
    SiLU, forget gate, then a weight-tied projection back to d.
  * learned persona vectors condition the style/role of the model.
  * rotary position embeddings, tied input/output embeddings.
"""

import math
from functools import lru_cache

# Chunk size for the log-space liquid scan. The scan renormalizes each chunk by
# exp(g_rel - m), so the chunk must satisfy  chunk * |log(gate_min)| < 709
# (float64 exp overflow threshold). Gates are clamped to >= 1e-12, i.e. max
# per-step decay 27.63; chunk 16 gives max exp argument 442 -- provably safe.
SCAN_CHUNK = 16

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

from .config import TinyLiquidConfig


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

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


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


@lru_cache(maxsize=8)
def _rope_freqs(seq_len: int, dim: int, theta: float, device: str, dtype: torch.dtype):
    half = dim // 2
    inv_freq = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
    t = torch.arange(seq_len, device=device, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)                       # (seq, half)
    cos = freqs.cos().to(dtype)
    sin = freqs.sin().to(dtype)
    return cos, sin


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    x = x.float()
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    x_rope = torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
    return x_rope.to(x.dtype if hasattr(x, "dtype") else torch.float32)


class BasisExpansion(nn.Module):
    """Liquid-style expansion: hidden -> N*B, group-norm over B, SiLU,
    forget gate, weight-tied projection back to hidden."""

    def __init__(self, cfg: TinyLiquidConfig):
        super().__init__()
        d = cfg.d_model
        self.n, self.b = cfg.basis_n, cfg.basis_b
        self.expand = cfg.basis_n * cfg.basis_b
        # in and forget-gate weights; output projection reuses w (tying)
        self.w = nn.Parameter(torch.empty(self.expand, d))
        self.w_forget = nn.Parameter(torch.empty(self.expand, d))
        self.gn = nn.GroupNorm(self.n, self.expand)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.normal_(self.w, std=0.02 / math.sqrt(self.expand))
        nn.init.normal_(self.w_forget, std=0.02 / math.sqrt(self.expand))

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        xr = apply_rope(x, cos, sin)                       # (b, s, d)
        e = xr @ self.w.t()                                # (b, s, N*B)
        e = e.transpose(1, 2)                              # (b, N*B, s) for groupnorm
        e = F.silu(self.gn(e))
        e = e.transpose(1, 2)
        f = torch.sigmoid(xr @ self.w_forget.t())          # forget/decay gate
        # Causal liquid recurrence: state_t = f_t * state_{t-1} + e_t.
        # Chunked log-space scan: exact math, bounded range per chunk, no
        # catastrophic cancellation, and far fewer Python iterations.
        G = torch.cumsum(torch.log(f.clamp_min(1e-12)), dim=1).double()
        b, s, E = e.shape
        h = torch.empty_like(e)
        state = torch.zeros(b, E, dtype=torch.float64)
        chunk = SCAN_CHUNK
        for start in range(0, s, chunk):
            end = min(start + chunk, s)
            base = G[:, start - 1:start] if start > 0 else G[:, :1]
            g_rel = G[:, start:end] - base                 # <= 0, non-increasing
            m = g_rel[:, -1:]                              # min within chunk
            # Shift exponents by the chunk min so every exp() argument <= 0:
            # fully stable for any gate saturation (no exp overflow).
            S = torch.cumsum(e[:, start:end].double() * torch.exp(-(g_rel - m)), dim=1)
            hc = torch.exp(g_rel - m) * (state.unsqueeze(1) * torch.exp(m) + S)
            h[:, start:end] = hc.float()
            state = hc[:, -1]
        return h @ self.w                                  # weight-tied projection


class GatedMLP(nn.Module):
    """Gated MLP with sigmoid forget gate (dense)."""

    def __init__(self, d: int, h: int):
        super().__init__()
        self.up = nn.Linear(d, h, bias=False)
        self.gate = nn.Linear(d, h, bias=False)
        self.forget = nn.Linear(d, h, bias=False)
        self.down = nn.Linear(h, d, bias=False)
        self.reset_parameters()

    def reset_parameters(self):
        for w in (self.up, self.gate, self.forget):
            nn.init.normal_(w.weight, std=0.02 / math.sqrt(w.weight.shape[0]))
        nn.init.normal_(self.down.weight, std=0.02 / math.sqrt(self.down.weight.shape[1]))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.silu(self.gate(x)) * self.up(x)
        h = h * torch.sigmoid(self.forget(x))
        return self.down(h)


class ExpertMLP(GatedMLP):
    pass


class MoEMLP(nn.Module):
    """Mixture-of-experts gated MLP: top-k routing over small experts."""

    def __init__(self, cfg: TinyLiquidConfig):
        super().__init__()
        d = cfg.d_model
        h = cfg.expert_hidden or (cfg.mlp_ratio * d // 2)
        self.n_experts = cfg.num_experts
        self.k = cfg.num_experts_per_tok
        self.router = nn.Linear(d, cfg.num_experts, bias=False)
        self.experts = nn.ModuleList([ExpertMLP(d, h) for _ in range(cfg.num_experts)])
        nn.init.normal_(self.router.weight, std=0.02 / math.sqrt(d))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, s, d = x.shape
        logits = self.router(x).float()                   # (b, s, E)
        topk = torch.topk(logits, self.k, dim=-1)
        weights = F.softmax(topk.values, dim=-1)          # (b, s, k)
        flat = x.reshape(-1, d)                           # (b*s, d)
        idx = topk.indices.reshape(-1, self.k)            # (b*s, k)
        out = torch.zeros_like(flat)
        flat_weights = weights.reshape(-1, self.k)
        for j in range(self.k):
            e_idx = idx[:, j]                     # (b*s,)
            wj = flat_weights[:, j]               # (b*s,)
            for e in range(self.n_experts):
                mask = e_idx == e
                if mask.any():
                    out[mask] += wj[mask].unsqueeze(1) * self.experts[e](flat[mask])
        return out.view(b, s, d)


class LiquidBlock(nn.Module):
    def __init__(self, cfg: TinyLiquidConfig):
        super().__init__()
        d = cfg.d_model
        self.norm1 = RMSNorm(d, cfg.norm_eps)
        self.basis = BasisExpansion(cfg)
        self.norm2 = RMSNorm(d, cfg.norm_eps)
        if cfg.num_experts > 0:
            self.mlp = MoEMLP(cfg)
        else:
            self.mlp = GatedMLP(d, cfg.mlp_ratio * d)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        x = x + self.basis(self.norm1(x), cos, sin)
        x = x + self.mlp(self.norm2(x))
        return x


class TinyLiquid(nn.Module):
    def __init__(self, cfg: TinyLiquidConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.persona_emb = nn.Embedding(cfg.num_personas, cfg.d_model)
        self.blocks = nn.ModuleList([LiquidBlock(cfg) for _ in range(cfg.n_blocks)])
        self.norm_out = RMSNorm(cfg.d_model, cfg.norm_eps)
        if cfg.tie_embeddings:
            self.lm_head = None  # tied below
        else:
            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.tower = None
        if cfg.tower_d and cfg.tower_blocks:
            tc = TinyLiquidConfig(vocab_size=cfg.vocab_size, d_model=cfg.tower_d,
                                  basis_n=cfg.basis_n, basis_b=cfg.basis_b,
                                  mlp_ratio=cfg.mlp_ratio, num_personas=0,
                                  max_seq_len=cfg.max_seq_len, rope_theta=cfg.rope_theta)
            self.up_proj = nn.Parameter(torch.zeros(cfg.tower_d, cfg.d_model))
            self.down_proj = nn.Parameter(torch.zeros(cfg.d_model, cfg.tower_d))
            with torch.no_grad():
                for i in range(min(cfg.d_model, cfg.tower_d)):
                    self.up_proj[i, i] = 1.0  # identity for trunk dims, zero for new dims
            self.tower = nn.ModuleList([LiquidBlock(tc) for _ in range(cfg.tower_blocks)])
            self._identity_tower()
        if getattr(cfg, "mtp_heads", 0):
            self.mtp_heads = nn.ModuleList([
                nn.Sequential(nn.Linear(cfg.d_model, cfg.d_model), nn.SiLU())
                for _ in range(cfg.mtp_heads)])
        else:
            self.mtp_heads = None
        self.reset_parameters()

    def _identity_tower(self):
        """Tower blocks start as exact identity (baseline output unchanged)."""
        with torch.no_grad():
            for blk in self.tower:
                blk.basis.w.zero_(); blk.basis.w_forget.zero_()
                blk.basis.gn.weight.fill_(1.0); blk.basis.gn.bias.zero_()
                blk.mlp.up.weight.zero_(); blk.mlp.gate.weight.zero_()
                blk.mlp.forget.weight.zero_(); blk.mlp.down.weight.zero_()

    def reset_parameters(self):
        nn.init.normal_(self.tok_emb.weight, std=0.02)
        nn.init.normal_(self.persona_emb.weight, std=0.02)

    def forward(
        self,
        ids: torch.Tensor,
        persona_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        cfg = self.cfg
        x = self.tok_emb(ids)
        if persona_ids is not None:
            x = x + self.persona_emb(persona_ids).unsqueeze(1)
        seq = ids.shape[1]
        theta = cfg.rope_theta
        cos, sin = _rope_freqs(seq, cfg.d_model, theta, str(ids.device), x.dtype)
        for blk in self.blocks:
            x = blk(x, cos, sin)
        if self.tower is not None:
            cos_t, sin_t = _rope_freqs(seq, cfg.tower_d, theta, str(ids.device), x.dtype)
            t = x @ self.up_proj.t()                 # (b, s, tower_d)
            for tb in self.tower:
                t = tb(t, cos_t, sin_t)
            x = x + t @ self.down_proj.t()           # zero-init residual: baseline preserved
        x = self.norm_out(x)
        if cfg.tie_embeddings:
            logits = x @ self.tok_emb.weight.t()
        else:
            logits = self.lm_head(x)
        return logits

    def hidden(self, ids: torch.Tensor, persona_ids: torch.Tensor | None = None) -> torch.Tensor:
        """Final hidden states (b, s, d) after norm_out, tower included."""
        cfg = self.cfg
        x = self.tok_emb(ids)
        if persona_ids is not None:
            x = x + self.persona_emb(persona_ids).unsqueeze(1)
        seq = ids.shape[1]
        theta = cfg.rope_theta
        cos, sin = _rope_freqs(seq, cfg.d_model, theta, str(ids.device), x.dtype)
        for blk in self.blocks:
            x = blk(x, cos, sin)
        if self.tower is not None:
            cos_t, sin_t = _rope_freqs(seq, cfg.tower_d, theta, str(ids.device), x.dtype)
            t = x @ self.up_proj.t()
            for tb in self.tower:
                t = tb(t, cos_t, sin_t)
            x = x + t @ self.down_proj.t()
        return self.norm_out(x)

    def forward_mtp(self, ids: torch.Tensor,
                    persona_ids: torch.Tensor | None = None):
        """Main logits + aux logits for multi-token prediction (Meta MTP).

        Each aux head predicts tokens at offset +2..+N+1 with a SiLU MLP whose
        output is projected by the TIED embedding (no new vocab-sized params).
        Returns (logits, [aux_logits_k]).
        """
        x = self.hidden(ids, persona_ids)
        cfg = self.cfg
        if cfg.tie_embeddings:
            logits = x @ self.tok_emb.weight.t()
        else:
            logits = self.lm_head(x)
        aux = []
        if self.mtp_heads is not None:
            for head in self.mtp_heads:
                aux.append(head(x) @ self.tok_emb.weight.t())
        return logits, aux

    @torch.no_grad()
    def encode(self, ids: torch.Tensor, persona_ids: torch.Tensor | None = None) -> torch.Tensor:
        """Final hidden states (b, s, d) after norm_out; no LM head."""
        return self.hidden(ids, persona_ids)

    def num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

    @torch.no_grad()
    def generate(
        self,
        tokenizer,
        prompt_ids,
        persona_id=0,
        max_new=200,
        temperature=0.8,
        top_k=40,
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        on_token=None,
    ):
        self.eval()
        ids = torch.tensor([prompt_ids], dtype=torch.long)
        stop_ids = {
            tok_id for tok_id in (
                tokenizer.token_to_id("<|endoftext|>"),
                tokenizer.token_to_id("<|user|>"),
                tokenizer.token_to_id("<|assistant|>"),
            )
            if tok_id is not None
        }
        for _ in range(max_new):
            window = ids[:, -self.cfg.max_seq_len :]
            logits = self(window, persona_ids=torch.tensor([persona_id]) if persona_id else None)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if repetition_penalty > 1.0 and ids.shape[1] > 8:
                seen = ids[0, -64:].unique()
                logits[:, seen] /= repetition_penalty
            if no_repeat_ngram_size > 0 and ids.shape[1] >= no_repeat_ngram_size:
                seq = ids[0].tolist()
                n = no_repeat_ngram_size
                prefix = tuple(seq[-(n - 1):])
                banned = set()
                for i in range(len(seq) - n + 1):
                    if tuple(seq[i:i + n - 1]) == prefix:
                        banned.add(seq[i + n - 1])
                if banned:
                    logits[:, list(banned)] = -float("inf")
            if top_k > 0:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, -1:]] = -float("inf")
            probs = F.softmax(logits.float(), dim=-1)
            nxt = torch.multinomial(probs, 1)
            ids = torch.cat([ids, nxt], dim=1)
            nxt_id = int(nxt.item())
            if on_token is not None:
                on_token(nxt_id)
            if nxt_id in stop_ids:
                break
        return ids[0].tolist()