"""Meiosis: PICO release 1 (2026-07). Tied-embedding looped decoder-only LM.

Spec: research/2026-07-first-release/final-spec.md (approved 2026-07-02).
embed -> prelude x1 -> [body of `body_blocks` distinct blocks xK loops,
per-loop LoRA + loop embed, Deep Delta vdim1 residuals] -> coda x1
-> RMSNorm -> tied unembed. Attention is MHA by default, GQA when
`n_kv_heads` < `n_heads` (2026-07-05 overhaul knobs, ADR-0009).
"""

import math
from dataclasses import dataclass

import torch
from torch import Tensor, nn
from torch.nn import functional

EMBED_STD = 0.02
LOOP_EMBED_STD = 0.02


@dataclass
class MeiosisConfig:
    # defaults = release shape per ADR-0009 (B'-GQA overhaul, 2026-07-05):
    # 3-block GQA body x3 loops, vocab 4096, ~5.76M total under the <6M cap
    vocab_size: int = 4096
    dim: int = 288
    n_heads: int = 6
    n_kv_heads: int | None = 2  # None -> MHA (= n_heads)
    ffn_hidden: int = 768
    prelude_layers: int = 1
    coda_layers: int = 1
    body_blocks: int = 3  # distinct blocks in the loop body
    max_loops: int = 4
    train_loops: int = 3
    lora_rank: int = 16
    rope_base: float = 10000.0
    max_seq_len: int = 512
    ddl_beta_init: float = 1.0
    # rms_norm backward amplifies grads by 1/sqrt(eps_rms) when k_in ~ 0 — which is
    # exactly the zero-init state. 1e-5 gave a 1.7e6x amplifier (1e5-magnitude grad
    # spikes; > fp16 max at ANY loss scale — the 2026-07-06 fp16 divergence, ADR-0012).
    # 1e-2 caps it at 1.7e3: fp16-safe, and identical bf16 training curves.
    ddl_k_eps: float = 1e-2
    ddl_v_sigmoid_scale: float = 4.0
    # intra-document attention (ADR-0019): tokens attend only within their own
    # EOS-delimited document. None = plain causal (pre-mask checkpoints).
    doc_mask_eos: int | None = 2

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

    @property
    def kv_heads(self) -> int:
        return self.n_kv_heads if self.n_kv_heads is not None else self.n_heads

    @property
    def qkv_dim(self) -> int:
        return self.dim + 2 * self.kv_heads * self.head_dim


def build_rope_cache(config: MeiosisConfig, length: int) -> tuple[Tensor, Tensor]:
    positions = torch.arange(length, dtype=torch.float32)
    inv_freq = 1.0 / (
        config.rope_base
        ** (torch.arange(0, config.head_dim, 2, dtype=torch.float32) / config.head_dim)
    )
    angles = torch.outer(positions, inv_freq)
    return torch.cos(angles), torch.sin(angles)


def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    rotated_even = x_even * cos - x_odd * sin
    rotated_odd = x_even * sin + x_odd * cos
    return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2)


def build_doc_mask(tokens: Tensor, eos_id: int) -> Tensor:
    """(B,T) tokens -> (B,1,T,T) bool, True where attention is allowed:
    causal AND same document. Exclusive EOS scan, so an EOS token is the
    last token of its document (FSX-1 convention)."""
    is_eos = tokens == eos_id
    doc_id = torch.cumsum(is_eos, dim=1) - is_eos.to(torch.long)
    same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
    causal = torch.ones(
        tokens.shape[1], tokens.shape[1], dtype=torch.bool, device=tokens.device
    ).tril()
    return (same & causal).unsqueeze(1)


class SwiGlu(nn.Module):
    def __init__(self, dim: int, hidden: int) -> None:
        super().__init__()
        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        gate, up = self.gate_up(x).chunk(2, dim=-1)
        return self.down(functional.silu(gate) * up)


class Attention(nn.Module):
    """MHA, or GQA when kv_heads < n_heads (KV repeated to full head count)."""

    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.n_heads = config.n_heads
        self.kv_heads = config.kv_heads
        self.head_dim = config.head_dim
        self.qkv = nn.Linear(config.dim, config.qkv_dim, bias=False)
        self.out = nn.Linear(config.dim, config.dim, bias=False)

    def forward(
        self,
        x: Tensor,
        cos: Tensor,
        sin: Tensor,
        qkv_delta: Tensor | None = None,
        attn_mask: Tensor | None = None,
    ) -> Tensor:
        batch, seq_len, dim = x.shape
        kv_dim = self.kv_heads * self.head_dim
        qkv = self.qkv(x)
        if qkv_delta is not None:
            qkv = qkv + qkv_delta
        q, k, v = qkv.split([dim, kv_dim, kv_dim], dim=-1)
        q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, seq_len, self.kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, seq_len, self.kv_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if self.kv_heads != self.n_heads:
            k = k.repeat_interleave(self.n_heads // self.kv_heads, dim=1)
            v = v.repeat_interleave(self.n_heads // self.kv_heads, dim=1)
        if attn_mask is not None:
            attended = functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        else:
            attended = functional.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(attended.transpose(1, 2).reshape(batch, seq_len, dim))


class DeepDeltaResidual(nn.Module):
    """DDL vdim1 (arXiv 2601.00417): x <- x + beta * (v - k^T x) * k.

    k is the sublayer output (rms-normed), beta in [0,2] gates between
    identity / projection / reflection, v is a learned scalar target.
    Replaces the plain additive residual in the looped block only.
    """

    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.k_eps = config.ddl_k_eps
        self.v_sigmoid_scale = config.ddl_v_sigmoid_scale
        self.beta = nn.Linear(config.dim, 1, bias=True)
        self.v_proj = nn.Linear(config.dim, 1, bias=True)
        beta_p = min(max(config.ddl_beta_init, 0.0), 2.0) / 2.0
        beta_p = min(max(beta_p, 1e-6), 1.0 - 1e-6)
        with torch.no_grad():
            self.beta.bias.fill_(math.log(beta_p) - math.log(1.0 - beta_p))

    def forward(self, x: Tensor, *, k_in: Tensor, context: Tensor) -> Tensor:
        k_dim = k_in.size(-1)
        eps_rms = (self.k_eps * self.k_eps) / k_dim
        k_rms = functional.rms_norm(k_in, [k_dim], eps=eps_rms)
        k_scale = 1.0 / math.sqrt(k_dim)
        beta = 2.0 * torch.sigmoid(self.beta(context).float())
        proj = torch.sum(k_rms * x, dim=-1, keepdim=True, dtype=torch.float32) * k_scale
        v = torch.sigmoid(self.v_proj(x).float()) * self.v_sigmoid_scale
        delta = ((beta * (v - proj)) * k_scale).to(dtype=x.dtype)
        return x + delta * k_rms


class Block(nn.Module):
    """Pre-norm block with plain additive residuals (prelude/coda)."""

    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.attn_norm = nn.RMSNorm(config.dim)
        self.attn = Attention(config)
        self.ffn_norm = nn.RMSNorm(config.dim)
        self.ffn = SwiGlu(config.dim, config.ffn_hidden)

    def forward(
        self, x: Tensor, cos: Tensor, sin: Tensor, attn_mask: Tensor | None = None
    ) -> Tensor:
        x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask=attn_mask)
        return x + self.ffn(self.ffn_norm(x))


class LoopedBlock(nn.Module):
    """Shared block with Deep Delta residuals; run K times with per-loop LoRA."""

    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.attn_norm = nn.RMSNorm(config.dim)
        self.attn = Attention(config)
        self.ddl_attn = DeepDeltaResidual(config)
        self.ffn_norm = nn.RMSNorm(config.dim)
        self.ffn = SwiGlu(config.dim, config.ffn_hidden)
        self.ddl_ffn = DeepDeltaResidual(config)

    def forward(
        self,
        x: Tensor,
        cos: Tensor,
        sin: Tensor,
        qkv_delta: Tensor,
        loop_emb: Tensor,
        attn_mask: Tensor | None = None,
    ) -> Tensor:
        # loop_emb conditions the sublayer inputs only — it is not carried in
        # the residual stream, so the block stays exactly identity at init
        # (zero-init out-projections -> k=0 -> DDL no-op) for any loop count.
        x_norm = self.attn_norm(x + loop_emb)
        x = self.ddl_attn(
            x,
            k_in=self.attn(x_norm, cos, sin, qkv_delta, attn_mask=attn_mask),
            context=x_norm,
        )
        x_norm = self.ffn_norm(x + loop_emb)
        return self.ddl_ffn(x, k_in=self.ffn(x_norm), context=x_norm)


class LoopLora(nn.Module):
    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.down = nn.ModuleList(
            nn.Linear(config.dim, config.lora_rank, bias=False)
            for _ in range(config.max_loops)
        )
        self.up = nn.ModuleList(
            nn.Linear(config.lora_rank, config.qkv_dim, bias=False)
            for _ in range(config.max_loops)
        )
        for up in self.up:
            nn.init.zeros_(up.weight)

    def forward(self, x: Tensor, loop_index: int) -> Tensor:
        clamped = min(loop_index, len(self.down) - 1)
        return self.up[clamped](self.down[clamped](x))


class Meiosis(nn.Module):
    def __init__(self, config: MeiosisConfig) -> None:
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.dim)
        self.prelude = nn.ModuleList(Block(config) for _ in range(config.prelude_layers))
        self.body = nn.ModuleList(LoopedBlock(config) for _ in range(config.body_blocks))
        self.loop_lora = nn.ModuleList(LoopLora(config) for _ in range(config.body_blocks))
        self.loop_embed = nn.Embedding(config.max_loops, config.dim)
        self.coda = nn.ModuleList(Block(config) for _ in range(config.coda_layers))
        self.final_norm = nn.RMSNorm(config.dim)
        cos, sin = build_rope_cache(config, config.max_seq_len)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        self.register_buffer("last_loop_rms", torch.zeros(config.max_loops), persistent=False)

    def forward(
        self,
        tokens: Tensor,
        loops: int | None = None,
        return_hidden: bool = False,
        collect_loop_rms: bool = False,
        attn_mask: Tensor | None = None,
    ) -> Tensor | tuple[Tensor, Tensor]:
        loop_count = loops if loops is not None else self.config.train_loops
        seq_len = tokens.shape[1]
        if seq_len > self.config.max_seq_len:
            raise ValueError(f"seq_len {seq_len} > max {self.config.max_seq_len}")
        x = self.embed(tokens)
        if attn_mask is None and self.config.doc_mask_eos is not None:
            attn_mask = build_doc_mask(tokens, self.config.doc_mask_eos)
        device_type = tokens.device.type
        compute_dtype = (
            torch.get_autocast_dtype(device_type)
            if torch.is_autocast_enabled(device_type)
            else x.dtype
        )
        cos = self.rope_cos[:seq_len].to(compute_dtype)
        sin = self.rope_sin[:seq_len].to(compute_dtype)
        for block in self.prelude:
            x = block(x, cos, sin, attn_mask=attn_mask)
        rms_per_loop = []
        for i in range(loop_count):
            clamped = min(i, self.config.max_loops - 1)
            loop_emb = self.loop_embed.weight[clamped]
            for block, lora in zip(self.body, self.loop_lora):
                x = block(x, cos, sin, lora(x + loop_emb, i), loop_emb, attn_mask=attn_mask)
            rms = x.float().pow(2).mean().sqrt()
            self.last_loop_rms[clamped] = rms.detach()
            if collect_loop_rms:
                rms_per_loop.append(rms)
        for block in self.coda:
            x = block(x, cos, sin, attn_mask=attn_mask)
        x = self.final_norm(x)
        out = x if return_hidden else functional.linear(x, self.embed.weight)
        if collect_loop_rms:
            return out, torch.stack(rms_per_loop)
        return out


def init_meiosis(model: Meiosis) -> None:
    """Mandatory MythosMini-validated init. Never mu-center the tied embedding."""
    with torch.no_grad():
        model.embed.weight.normal_(mean=0.0, std=EMBED_STD)
        model.loop_embed.weight.normal_(mean=0.0, std=LOOP_EMBED_STD)
        for block in [*model.prelude, *model.body, *model.coda]:
            nn.init.zeros_(block.attn.out.weight)
            nn.init.zeros_(block.ffn.down.weight)


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def muon_param_split(model: Meiosis) -> tuple[list[nn.Parameter], list[nn.Parameter]]:
    """Explicit Muon/aux split (ADR-0005). Muon gets the block and LoRA
    matrices; the tied embedding, loop embeddings, norm gains, and 1-row DDL
    heads stay on NAdamW. Listed explicitly - no shape heuristics, so a
    rank-8 pilot LoRA cannot silently fall out of the Muon group.
    """
    muon: list[nn.Parameter] = []
    for block in [*model.prelude, *model.body, *model.coda]:
        muon += [
            block.attn.qkv.weight,
            block.attn.out.weight,
            block.ffn.gate_up.weight,
            block.ffn.down.weight,
        ]
    for lora in model.loop_lora:
        muon += [linear.weight for linear in [*lora.down, *lora.up]]
    muon_ids = {id(p) for p in muon}
    aux = [p for p in model.parameters() if id(p) not in muon_ids]
    return muon, aux