"""Sol Pro 138M: a width/depth-scaled Sol Flash causal language model."""

from __future__ import annotations

from dataclasses import dataclass
import math

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


@dataclass(frozen=True)
class SolProConfig:
    vocab_size: int = 4096
    width: int = 768
    heads: int = 24
    kv_heads: int = 8
    stored_blocks: int = 12
    ffn_width: int = 3968
    recurrent_start: int = 1
    recurrent_blocks: int = 4
    recurrent_passes: int = 2
    tn_gram_rank: int = 299
    max_ngram_order: int = 4
    use_tn_gram: 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

    @property
    def effective_blocks(self) -> int:
        return self.stored_blocks + self.recurrent_blocks * (self.recurrent_passes - 1)

    def validate(self) -> None:
        if self.width % self.heads:
            raise ValueError("width must be divisible by heads")
        if self.heads % self.kv_heads:
            raise ValueError("heads must be divisible by kv_heads")
        if self.head_dim % 2:
            raise ValueError("head_dim must be even for RoPE")
        if self.recurrent_start + self.recurrent_blocks > self.stored_blocks:
            raise ValueError("recurrent block range exceeds stored blocks")
        if self.max_ngram_order != 4:
            raise ValueError("Sol Pro TN-gram currently supports orders 2 through 4")
        if self.tn_gram_rank < 1:
            raise ValueError("tn_gram_rank must be positive")


def model_config() -> SolProConfig:
    return SolProConfig()


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:
        normalized = x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps)
        return normalized.to(x.dtype) * self.weight.to(x.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)


def _linear_2d(module: nn.Module, hidden: torch.Tensor) -> torch.Tensor:
    shape = hidden.shape
    projected = module(hidden.reshape(-1, shape[-1]))
    return projected.reshape(*shape[:-1], -1)


class TNGramMemory(nn.Module):
    """Shared tensorized 2-, 3-, and 4-gram memory with CP-style factors."""

    def __init__(self, cfg: SolProConfig):
        super().__init__()
        self.vocab_size = cfg.vocab_size
        self.rank = cfg.tn_gram_rank
        self.width = cfg.width
        self.position_factors = nn.ModuleList(
            [nn.Embedding(cfg.vocab_size, self.rank) for _ in range(cfg.max_ngram_order)]
        )
        # Shared vectors absorb the unused oldest positions for lower orders.
        self.order_absorption = nn.Parameter(torch.ones(cfg.max_ngram_order - 2, self.rank))
        self.log_order_scales = nn.Parameter(torch.zeros(cfg.max_ngram_order - 1))
        self.key_projection = nn.Linear(3 * self.rank, cfg.width, bias=False)
        self.value_projection = nn.Linear(3 * self.rank, cfg.width, bias=False)
        self.query_norm = RMSNorm(cfg.width)
        self.key_norm = RMSNorm(cfg.width)
        self.value_norm = RMSNorm(cfg.width)
        self.local_conv = nn.Conv1d(
            cfg.width,
            cfg.width,
            kernel_size=3,
            dilation=cfg.max_ngram_order,
            groups=cfg.width,
            bias=True,
        )
        self.residual_scale = nn.Parameter(torch.tensor(0.1))

    @staticmethod
    def _lag(ids: torch.Tensor, lag: int) -> torch.Tensor:
        if lag == 0:
            return ids
        return F.pad(ids, (lag, 0), value=0)[:, : ids.shape[1]]

    @staticmethod
    def _rms_normalize(x: torch.Tensor) -> torch.Tensor:
        normalized = x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
        return normalized.to(x.dtype)

    def _order_features(self, ids: torch.Tensor) -> torch.Tensor:
        # Factors A1..A4 correspond to oldest-to-newest token positions.
        by_position = [
            self.position_factors[index](self._lag(ids, 3 - index))
            for index in range(4)
        ]
        absorption = self.order_absorption.to(by_position[0].dtype)
        bigram = by_position[2] * by_position[3]
        bigram = bigram * absorption[0] * absorption[1]
        trigram = by_position[1] * by_position[2] * by_position[3]
        trigram = trigram * absorption[0]
        quagram = by_position[0] * by_position[1] * by_position[2] * by_position[3]

        orders = [bigram, trigram, quagram]
        scales = self.log_order_scales.exp().to(by_position[0].dtype)
        normalized = [
            self._rms_normalize(features) * scales[index]
            for index, features in enumerate(orders)
        ]
        return torch.cat(normalized, dim=-1)

    def forward(self, ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
        features = self._order_features(ids)
        query = self.query_norm(hidden)
        key = self.key_norm(_linear_2d(self.key_projection, features))
        value = _linear_2d(self.value_projection, features)
        similarity = (query * key).sum(dim=-1, keepdim=True) / math.sqrt(self.width)
        gate_logits = similarity.sign() * similarity.abs().clamp_min(1e-6).sqrt()
        gated_value = torch.sigmoid(gate_logits).to(value.dtype) * value

        conv_input = self.value_norm(gated_value).transpose(1, 2)
        causal_padding = (self.local_conv.kernel_size[0] - 1) * self.local_conv.dilation[0]
        conv_input = F.pad(conv_input, (causal_padding, 0))
        local_update = F.silu(self.local_conv(conv_input).transpose(1, 2))
        memory = gated_value + local_update
        return hidden + self.residual_scale.tanh().to(hidden.dtype) * memory.to(hidden.dtype)


class XSAAttention(nn.Module):
    """Sol's causal GQA with per-head Q/K RMSNorm and XSA value subtraction."""

    def __init__(self, cfg: SolProConfig):
        super().__init__()
        self.cfg = cfg
        self.qkv = nn.Linear(
            cfg.width,
            (cfg.heads + 2 * 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_width = self.cfg.heads * self.cfg.head_dim
        kv_width = self.cfg.kv_heads * self.cfg.head_dim
        q_flat, k_flat, v_flat = _linear_2d(self.qkv, x).split(
            (q_width, kv_width, kv_width), dim=-1
        )
        q = q_flat.view(batch, length, self.cfg.heads, self.cfg.head_dim).transpose(1, 2)
        k = k_flat.view(batch, length, self.cfg.kv_heads, self.cfg.head_dim).transpose(1, 2)
        v = v_flat.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
        attended = F.scaled_dot_product_attention(
            q, k, v, is_causal=True, enable_gqa=(groups > 1)
        )
        if groups > 1:
            unit_v = F.normalize(v, p=2, dim=-1, eps=1e-6).unsqueeze(2)
            grouped = attended.view(
                batch, self.cfg.kv_heads, groups, length, self.cfg.head_dim
            )
            attended = (
                grouped - (grouped * unit_v).sum(-1, keepdim=True) * unit_v
            ).reshape(batch, self.cfg.heads, length, self.cfg.head_dim)
        else:
            unit_v = F.normalize(v, p=2, dim=-1, eps=1e-6)
            attended = attended - (attended * unit_v).sum(-1, keepdim=True) * unit_v
        attended = attended.transpose(1, 2).contiguous().view(batch, length, -1)
        return _linear_2d(self.o, attended)


class Block(nn.Module):
    def __init__(self, cfg: SolProConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.width)
        self.attn = XSAAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.width)
        self.gate_up = nn.Linear(cfg.width, 2 * 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)
        gate_up = _linear_2d(self.gate_up, normed)
        gate, up = gate_up.chunk(2, dim=-1)
        return x + _linear_2d(self.down, F.silu(gate) * up)


class SolProForCausalLM(nn.Module):
    """Sol Flash architecture scaled to a 138,037,594-parameter model."""

    def __init__(self, cfg: SolProConfig | None = None):
        super().__init__()
        self.config = cfg or SolProConfig()
        self.config.validate()
        cfg = self.config
        self.embedding = nn.Embedding(cfg.vocab_size, cfg.width)
        self.tn_gram = TNGramMemory(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.embedding.weight, std=0.02)
        if self.tn_gram is not None:
            for factor in self.tn_gram.position_factors:
                nn.init.normal_(factor.weight, std=1.0)
        nn.init.normal_(self.loop_embeddings, std=0.01)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        if ids.shape[1] > self.config.max_position_embeddings:
            raise ValueError("input sequence exceeds configured context length")
        x = self.embedding(ids)
        if torch.is_autocast_enabled("cuda"):
            x = x.to(torch.get_autocast_dtype("cuda"))
        if 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].to(x.dtype)
            for local_index, block in enumerate(self.blocks[start:stop]):
                proposal = block(x)
                if self.config.use_loop_conditioning:
                    gate = torch.sigmoid(self.loop_gates[pass_index, local_index]).to(x.dtype)[None, None, :]
                    x = x + gate * (proposal - x)
                else:
                    x = proposal
        for block in self.blocks[stop:]:
            x = block(x)
        return F.linear(self.norm(x), self.embedding.weight)


class SolProCausalLoss(nn.Module):
    """Return BF16 cross-entropy directly to keep the training graph compact."""

    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model
        self.vocab_size = model.config.vocab_size

    def forward(self, token_batch: torch.Tensor) -> torch.Tensor:
        logits = self.model(token_batch[:, :-1])
        targets = token_batch[:, 1:]
        return F.cross_entropy(
            logits.reshape(-1, self.vocab_size), targets.reshape(-1)
        )


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

def load_model(model_dir, device=None):
    """Load the original tokenizer and released Sol Pro inference weights."""
    import json
    from pathlib import Path
    from safetensors.torch import load_file
    from tokenizers import Tokenizer
    root = Path(model_dir)
    settings = json.loads((root / "config.json").read_text())
    settings = {name: value for name, value in settings.items() if name in SolProConfig.__dataclass_fields__}
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
    model = SolProForCausalLM(SolProConfig(**settings)).to(device=device, dtype=dtype)
    model.load_state_dict(load_file(str(root / "model.safetensors"), device="cpu"), strict=True)
    model.eval()
    tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
    return model, tokenizer

@torch.inference_mode()
def generate(model, tokenizer, prompt, max_new_tokens=64, temperature=0.0):
    """Return prompt plus continuation, without exceeding the context window."""
    if max_new_tokens < 0:
        raise ValueError("max_new_tokens must be nonnegative")
    ids = tokenizer.encode(prompt, add_special_tokens=False).ids
    eos = tokenizer.token_to_id("<eos>")
    if not ids:
        if eos is None:
            raise ValueError("An empty prompt requires an EOS token")
        ids = [eos]
    capacity = model.config.max_position_embeddings
    if len(ids) + max_new_tokens > capacity:
        raise ValueError("Prompt and requested continuation exceed the 2,048-token context")
    device = next(model.parameters()).device
    for _ in range(max_new_tokens):
        scores = model(torch.tensor([ids], device=device, dtype=torch.long))[0, -1].float()
        if temperature > 0:
            next_id = int(torch.multinomial(torch.softmax(scores / temperature, dim=-1), 1).item())
        else:
            next_id = int(scores.argmax().item())
        ids.append(next_id)
        if next_id == eos:
            break
    return tokenizer.decode(ids, skip_special_tokens=True)