"""Sol Lite source architecture and parameter-audited Sol Nano variant."""

from __future__ import annotations

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

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

try:
    from flash_attn.cute import flash_attn_func as _flash_attn4_func
except ImportError:
    _flash_attn4_func = None

ATTENTION_BACKEND = os.environ.get("SOL_NANO_ATTENTION", "triton")
if ATTENTION_BACKEND == "triton":
    from torch.nn.attention.flex_attention import flex_attention as _flex_attention

    def _causal_score(score, batch, head, query_index, key_index):
        return torch.where(query_index >= key_index, score, -float("inf"))

    @torch.compile(backend="inductor", fullgraph=True, dynamic=False)
    def _triton_causal_attention(q, k, v):
        return _flex_attention(q, k, v, score_mod=_causal_score, enable_gqa=True)


@torch.compiler.disable
def fused_causal_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    """Fused causal GQA on (batch, sequence, heads, head_dim), with no SDPA."""
    if ATTENTION_BACKEND == "triton":
        result = _triton_causal_attention(
            q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(),
            v.transpose(1, 2).contiguous(),
        )
        return result.transpose(1, 2)
    if _flash_attn4_func is None:
        raise RuntimeError(
            "Sol Nano requires flash-attn-4; install the pinned FA4 runtime dependency"
        )
    # SM120's current CuTe packed-GQA LSE path does not compile; unpacked GQA
    # preserves the same attention math and shared K/V projection parameters.
    result = _flash_attn4_func(q, k, v, causal=True, pack_gqa=False)
    # FA4 currently returns (output, logsumexp), even when LSE is not requested.
    return result[0] if isinstance(result, tuple) else result


@dataclass(frozen=True)
class SolLiteConfig:
    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_tn_gram: bool = False
    tn_rank: int = 10
    tn_buckets: int = 4950
    use_qk_norm: bool = True
    use_loop_conditioning: bool = True
    use_xsa: 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.to(x.dtype)


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: SolLiteConfig):
        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 TNGramLite(nn.Module):
    """Factorized causal 2--5-gram lookup with fixed-size hash tables."""

    def __init__(self, cfg: SolLiteConfig):
        super().__init__()
        self.rank = cfg.tn_rank
        self.buckets = cfg.tn_buckets
        self.token_factors = nn.Parameter(torch.empty(cfg.vocab_size, self.rank))
        self.hash_tables = nn.Parameter(torch.empty(4, self.buckets, self.rank))
        self.order_factors = nn.Parameter(torch.ones(4, self.rank))
        self.rank_to_hidden = nn.Parameter(torch.empty(cfg.width, self.rank))
        self.output_bias = nn.Parameter(torch.zeros(cfg.width))
        self.gates = nn.Parameter(torch.full((4,), -2.0))
        self.order_mix = nn.Parameter(torch.zeros(4, self.rank))
        self.context_controls = nn.Parameter(torch.zeros(16))
        nn.init.normal_(self.token_factors, mean=1.0, std=0.02)
        nn.init.normal_(self.hash_tables, std=0.02)
        nn.init.normal_(self.rank_to_hidden, std=0.02)
        nn.init.normal_(self.order_mix, std=0.01)

    def forward(self, ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
        batch, length = ids.shape
        positions = torch.arange(length, device=ids.device)
        result = hidden + self.output_bias.to(hidden.dtype)
        for order_index, order in enumerate(range(2, 6)):
            hashed = torch.zeros_like(ids)
            factors = torch.ones(
                batch, length, self.rank, device=ids.device, dtype=torch.float32
            )
            for lag in range(order):
                if lag == 0:
                    shifted = ids
                else:
                    shifted = F.pad(ids[:, :-lag], (lag, 0), value=0)
                hashed = (hashed * (131 + order_index * 6) + shifted) % self.buckets
                factors = factors * self.token_factors[shifted].float()

            rank = factors * self.hash_tables[order_index, hashed].float()
            control = self.context_controls[order_index * 4 : (order_index + 1) * 4]
            rank = rank * self.order_factors[order_index]
            rank = rank + self.order_mix[order_index] * control.tanh().mean()
            contribution = F.linear(rank, self.rank_to_hidden.float()).to(hidden.dtype)
            valid = (positions >= order - 1).to(hidden.dtype)[None, :, None]
            result = result + contribution * torch.sigmoid(self.gates[order_index]) * valid
        return result


class CausalGQAAttention(nn.Module):
    def __init__(self, cfg: SolLiteConfig):
        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 _attend(self, x: torch.Tensor) -> tuple[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)
        # Both fused backends keep the grouped-query projection parameters.
        attended = fused_causal_attention(
            q.transpose(1, 2).contiguous(),
            k.transpose(1, 2).contiguous(),
            v.transpose(1, 2).contiguous(),
        )
        return attended, v.transpose(1, 2).contiguous()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attended, _ = self._attend(x)
        batch, length, _ = x.shape
        return self.o(attended.contiguous().view(batch, length, -1))


class XSAAttention(CausalGQAAttention):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attended, v = self._attend(x)
        batch, length, _ = x.shape
        groups = self.cfg.heads // self.cfg.kv_heads
        if groups > 1:
            v = v.repeat_interleave(groups, dim=2)
        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.contiguous().view(batch, length, -1))


class Block(nn.Module):
    def __init__(self, cfg: SolLiteConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.width)
        self.attn = XSAAttention(cfg) if cfg.use_xsa else CausalGQAAttention(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)
        q, k, v = split(self.q(x)), split(self.k(x)), split(self.v(x))
        attended = flash_attention4_causal(q.contiguous(), k.contiguous(), v.contiguous())
        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.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 ReleasedSolLiteControl(nn.Module):
    """Faithful PyTorch control for the released 2,996,480-parameter graph."""

    def __init__(self):
        super().__init__()
        self.config = SolLiteConfig(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 SolForCausalLM(nn.Module):
    """Sol Lite with optional EngramLite or factorized TNGramLite memory."""

    def __init__(self, cfg: SolLiteConfig = SolLiteConfig()):
        super().__init__()
        self.config = cfg
        if cfg.use_engram and cfg.use_tn_gram:
            raise ValueError("select either EngramLite or TNGramLite, not both")
        self.embedding = nn.Embedding(cfg.vocab_size, cfg.width)
        self.engram = EngramLite(cfg) if cfg.use_engram else None
        self.tn_gram = TNGramLite(cfg) if cfg.use_tn_gram 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)
        elif self.tn_gram is not None:
            x = self.tn_gram(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) -> SolLiteConfig:
    if name == "sol_nano_2p9m_tn_gram":
        return SolLiteConfig(
            vocab_size=1024,
            width=128,
            heads=4,
            kv_heads=2,
            stored_blocks=10,
            ffn_width=536,
            recurrent_start=1,
            recurrent_blocks=4,
            recurrent_passes=2,
            engram_entries=0,
            use_engram=False,
            use_tn_gram=True,
            tn_rank=10,
            tn_buckets=4950,
            use_qk_norm=True,
            use_loop_conditioning=True,
            use_xsa=False,
            max_position_embeddings=512,
            rope_theta=20000.0,
        )
    if name == "control":
        return SolLiteConfig(kv_heads=4, ffn_width=540, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False)
    if name == "gqa":
        return SolLiteConfig(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 SolLiteConfig(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 SolLiteConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=True, use_loop_conditioning=True)
    if name == "full":
        return SolLiteConfig()
    raise ValueError(f"unknown architecture variant: {name}")


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


def architecture_audit() -> dict[str, int | float | bool]:
    results = {}
    for name in ("released_control", "control", "gqa", "gqa_qknorm", "loop_conditioned", "full", "sol_nano_2p9m_tn_gram"):
        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 = SolLiteConfig.__dataclass_fields__
    cfg = SolLiteConfig(**{key: value for key, value in raw.items() if key in fields})
    model = SolForCausalLM(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: SolForCausalLM,
    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("SOL_LITE_AUDIT=" + json.dumps(architecture_audit(), sort_keys=True))