"""Parameter-audited Nero-XS-2 candidates for PyTorch/XLA."""

from __future__ import annotations

from dataclasses import dataclass
import json
import math
from pathlib import Path

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


@dataclass(frozen=True)
class NeroConfig:
    vocab_size: int = 2048
    width: int = 128
    heads: int = 4
    kv_heads: int = 2
    stored_blocks: int = 10
    ffn_width: int = 531
    recurrent_start: int = 1
    recurrent_blocks: int = 4
    recurrent_passes: int = 2
    engram_entries: int = 768
    use_engram: bool = True
    use_qk_norm: bool = True
    use_loop_conditioning: bool = True
    max_position_embeddings: int = 2048
    rope_theta: float = 20000.0

    @property
    def head_dim(self) -> int:
        return self.width // self.heads


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight


class CenteredUnitNorm(nn.Module):
    def __init__(self, width: int, eps: float = 1e-5):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(width))
        self.shift = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        centered = x - x.mean(-1, keepdim=True)
        return centered * torch.rsqrt(centered.square().mean(-1, keepdim=True) + self.eps) * self.scale + self.shift


def deterministic_coordinates(length: int, width: int, base: float, device, dtype):
    half = (width + 1) // 2
    positions = torch.arange(length, device=device, dtype=torch.float32)[:, None]
    frequencies = torch.exp(torch.arange(half, device=device, dtype=torch.float32) * (-math.log(base) / max(half - 1, 1)))
    result = torch.cat((torch.sin(positions * frequencies), torch.cos(positions * frequencies)), dim=-1)[:, :width]
    return result.to(dtype)


def apply_rope(x: torch.Tensor, theta: float) -> torch.Tensor:
    _, _, length, dim = x.shape
    inv = theta ** (-torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim)
    angles = torch.arange(length, device=x.device, dtype=torch.float32)[:, None] * inv[None, :]
    cos = angles.cos().to(x.dtype)[None, None, :, :]
    sin = angles.sin().to(x.dtype)[None, None, :, :]
    even, odd = x[..., ::2], x[..., 1::2]
    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)


class EngramLite(nn.Module):
    """Collision-tolerant bigram/trigram memory with a contextual read gate."""

    def __init__(self, cfg: NeroConfig):
        super().__init__()
        self.entries = cfg.engram_entries
        self.tables = nn.ModuleList([nn.Embedding(cfg.engram_entries, cfg.width) for _ in range(2)])
        self.gate = nn.Linear(cfg.width, 2, bias=True)
        self.scale = nn.Parameter(torch.tensor(0.1))

    def _hash(self, ids: torch.Tensor, order: int, prime: int) -> torch.Tensor:
        padded = F.pad(ids, (order - 1, 0), value=0)
        value = torch.zeros_like(ids)
        for offset in range(order):
            value = (value * prime + padded[:, offset : offset + ids.shape[1]]) % self.entries
        return value

    def forward(self, ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
        bigram = self.tables[0](self._hash(ids, 2, 10007))
        trigram = self.tables[1](self._hash(ids, 3, 10009))
        weights = torch.sigmoid(self.gate(hidden))
        memory = weights[..., :1] * bigram + weights[..., 1:] * trigram
        return hidden + self.scale.tanh() * memory


class XSAAttention(nn.Module):
    def __init__(self, cfg: NeroConfig):
        super().__init__()
        self.cfg = cfg
        self.q = nn.Linear(cfg.width, cfg.heads * cfg.head_dim, bias=False)
        self.k = nn.Linear(cfg.width, cfg.kv_heads * cfg.head_dim, bias=False)
        self.v = nn.Linear(cfg.width, cfg.kv_heads * cfg.head_dim, bias=False)
        self.o = nn.Linear(cfg.heads * cfg.head_dim, cfg.width, bias=False)
        self.q_norm = RMSNorm(cfg.head_dim) if cfg.use_qk_norm else nn.Identity()
        self.k_norm = RMSNorm(cfg.head_dim) if cfg.use_qk_norm else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        q = self.q(x).view(batch, length, self.cfg.heads, self.cfg.head_dim).transpose(1, 2)
        k = self.k(x).view(batch, length, self.cfg.kv_heads, self.cfg.head_dim).transpose(1, 2)
        v = self.v(x).view(batch, length, self.cfg.kv_heads, self.cfg.head_dim).transpose(1, 2)
        q = apply_rope(self.q_norm(q), self.cfg.rope_theta)
        k = apply_rope(self.k_norm(k), self.cfg.rope_theta)
        groups = self.cfg.heads // self.cfg.kv_heads
        if groups > 1:
            k = k.repeat_interleave(groups, dim=1)
            v = v.repeat_interleave(groups, dim=1)
        attended = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        unit_v = F.normalize(v, p=2, dim=-1, eps=1e-6)
        attended = attended - (attended * unit_v).sum(-1, keepdim=True) * unit_v
        return self.o(attended.transpose(1, 2).contiguous().view(batch, length, -1))


class Block(nn.Module):
    def __init__(self, cfg: NeroConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.width)
        self.attn = XSAAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.width)
        self.gate = nn.Linear(cfg.width, cfg.ffn_width, bias=False)
        self.up = nn.Linear(cfg.width, cfg.ffn_width, bias=False)
        self.down = nn.Linear(cfg.ffn_width, cfg.width, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.attn_norm(x))
        normed = self.ffn_norm(x)
        return x + self.down(F.silu(self.gate(normed)) * self.up(normed))


class ReleasedXSAAttention(nn.Module):
    def __init__(self, width: int = 128, heads: int = 4):
        super().__init__()
        self.width, self.heads, self.head_dim = width, heads, width // heads
        self.q = nn.Linear(width, width, bias=False)
        self.k = nn.Linear(width, width, bias=False)
        self.v = nn.Linear(width, width, bias=False)
        self.o = nn.Linear(width, width, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        split = lambda value: value.view(batch, length, self.heads, self.head_dim).transpose(1, 2)
        q, k, v = split(self.q(x)), split(self.k(x)), split(self.v(x))
        attended = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        coefficient = (attended * v).sum(-1, keepdim=True) / v.square().sum(-1, keepdim=True).clamp_min(1e-6)
        attended = attended - coefficient * v
        return self.o(attended.transpose(1, 2).contiguous().view(batch, length, self.width))


class ReleasedBlock(nn.Module):
    def __init__(self, width: int = 128, ffn_width: int = 540):
        super().__init__()
        self.attn_norm = CenteredUnitNorm(width)
        self.attn = ReleasedXSAAttention(width)
        self.ffn_norm = CenteredUnitNorm(width)
        self.expand = nn.Linear(width, 2 * ffn_width, bias=False)
        self.contract = nn.Linear(ffn_width, width, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.attn_norm(x))
        content, gate = self.expand(self.ffn_norm(x)).chunk(2, dim=-1)
        return x + self.contract(F.silu(content) * torch.sigmoid(gate))


class ReleasedNeroXSControl(nn.Module):
    """Faithful PyTorch control for the released 2,996,480-parameter graph."""

    def __init__(self):
        super().__init__()
        self.config = NeroConfig(kv_heads=4, ffn_width=540, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False)
        self.embedding = nn.Embedding(2048, 128)
        self.blocks = nn.ModuleList([ReleasedBlock() for _ in range(10)])
        self.norm = CenteredUnitNorm(128)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(ids) + deterministic_coordinates(ids.shape[1], 128, 20000.0, ids.device, self.embedding.weight.dtype)[None]
        x = self.blocks[0](x)
        for _ in range(2):
            for block in self.blocks[1:5]:
                x = block(x)
        for block in self.blocks[5:]:
            x = block(x)
        return F.linear(self.norm(x), self.embedding.weight)


class NeroXSA2ForCausalLM(nn.Module):
    """XSA + selective recurrence + loop conditioning + EngramLite."""

    def __init__(self, cfg: NeroConfig = NeroConfig()):
        super().__init__()
        self.config = cfg
        self.embedding = nn.Embedding(cfg.vocab_size, cfg.width)
        self.engram = EngramLite(cfg) if cfg.use_engram else None
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.stored_blocks)])
        self.loop_embeddings = nn.Parameter(torch.zeros(cfg.recurrent_passes, cfg.width))
        self.loop_gates = nn.Parameter(torch.zeros(cfg.recurrent_passes, cfg.recurrent_blocks, cfg.width))
        self.norm = RMSNorm(cfg.width)
        nn.init.normal_(self.loop_embeddings, std=0.01)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(ids)
        if self.engram is not None:
            x = self.engram(ids, x)
        start = self.config.recurrent_start
        stop = start + self.config.recurrent_blocks
        for block in self.blocks[:start]:
            x = block(x)
        for pass_index in range(self.config.recurrent_passes):
            if self.config.use_loop_conditioning:
                x = x + self.loop_embeddings[pass_index]
            for local_index, block in enumerate(self.blocks[start:stop]):
                if self.config.use_loop_conditioning:
                    proposal = block(x)
                    gate = torch.sigmoid(self.loop_gates[pass_index, local_index])[None, None, :]
                    x = x + gate * (proposal - x)
                else:
                    x = block(x)
        for block in self.blocks[stop:]:
            x = block(x)
        return F.linear(self.norm(x), self.embedding.weight)


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


def variant_config(name: str) -> NeroConfig:
    if name == "control":
        return NeroConfig(kv_heads=4, ffn_width=540, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False)
    if name == "gqa":
        return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False)
    if name == "gqa_qknorm":
        return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=True, use_loop_conditioning=False)
    if name == "loop_conditioned":
        return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=True, use_loop_conditioning=True)
    if name == "full":
        return NeroConfig()
    raise ValueError(f"unknown architecture variant: {name}")


def build_variant(name: str) -> nn.Module:
    if name == "released_control":
        return ReleasedNeroXSControl()
    return NeroXSA2ForCausalLM(variant_config(name))


def architecture_audit() -> dict[str, int | float | bool]:
    results = {}
    for name in ("released_control", "control", "gqa", "gqa_qknorm", "loop_conditioned", "full"):
        model = build_variant(name)
        cfg = model.config
        count = parameter_count(model)
        if count >= 3_000_000:
            raise AssertionError(f"{name} exceeds the parameter cap: {count:,}")
        ids = torch.arange(64).remainder(cfg.vocab_size).view(1, -1)
        with torch.no_grad():
            logits = model(ids)
            changed = ids.clone()
            changed[:, 32:] = (changed[:, 32:] + 17) % cfg.vocab_size
            changed_logits = model(changed)
        prefix_error = float((logits[:, :32] - changed_logits[:, :32]).abs().max())
        if logits.shape != (1, 64, cfg.vocab_size) or prefix_error > 1e-5:
            raise AssertionError((name, logits.shape, prefix_error))
        results[name] = {"parameters": count, "causal_prefix_max_error": prefix_error}
    return results


def load_model(model_dir: str | Path, device: str | torch.device = "cpu"):
    """Load the released safetensors checkpoint and tokenizer."""
    from safetensors.torch import load_file
    from transformers import AutoTokenizer

    model_dir = Path(model_dir)
    raw = json.loads((model_dir / "config.json").read_text())
    fields = NeroConfig.__dataclass_fields__
    cfg = NeroConfig(**{key: value for key, value in raw.items() if key in fields})
    model = NeroXSA2ForCausalLM(cfg)
    model.load_state_dict(load_file(model_dir / "model.safetensors"), strict=True)
    model.to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    return model, tokenizer


@torch.inference_mode()
def generate(
    model: NeroXSA2ForCausalLM,
    tokenizer,
    prompt: str,
    max_new_tokens: int = 64,
    temperature: float = 0.8,
    top_p: float = 0.95,
    repetition_penalty: float = 1.1,
    seed: int = 7,
) -> str:
    """Simple deterministic-seed nucleus sampler using full-prefix recomputation."""
    device = model.embedding.weight.device
    ids = tokenizer.encode(prompt, add_special_tokens=False)
    generator = torch.Generator(device=device).manual_seed(seed)
    for _ in range(max_new_tokens):
        context = torch.tensor([ids[-model.config.max_position_embeddings :]], device=device)
        logits = model(context)[0, -1].float()
        if repetition_penalty != 1.0:
            seen = torch.tensor(sorted(set(ids)), device=device)
            logits[seen] = torch.where(
                logits[seen] < 0,
                logits[seen] * repetition_penalty,
                logits[seen] / repetition_penalty,
            )
        if temperature <= 0:
            next_id = int(logits.argmax())
        else:
            sorted_logits, sorted_ids = (logits / temperature).sort(descending=True)
            probabilities = sorted_logits.softmax(-1)
            keep = probabilities.cumsum(-1) <= top_p
            keep[0] = True
            filtered = probabilities * keep
            choice = torch.multinomial(filtered / filtered.sum(), 1, generator=generator)
            next_id = int(sorted_ids[choice])
        ids.append(next_id)
        if next_id == tokenizer.eos_token_id:
            break
    return tokenizer.decode(ids, skip_special_tokens=True)


if __name__ == "__main__":
    import json
    print("NERO_XS2_AUDIT=" + json.dumps(architecture_audit(), sort_keys=True))