"""glint-router-1m: standalone inference. torch + tokenizers + safetensors, nothing else.

one forward pass over the prompt gives you every field at once: domain,
complexity, code, math, reasoning, long_output, route, and a 64-d projection
used for user-defined categories. no decoding loop, no output parsing.
"""

from __future__ import annotations

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

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

HERE = Path(__file__).parent
PAD_ID = 0
BOS_ID = 1
BINARY_FIELDS = ("code", "math", "reasoning", "long_output")
COMPLEXITY_LEVELS = 5

DOMAINS = (
    "programming", "web_dev", "databases", "devops_sysadmin", "security",
    "data_science_ml", "math", "science", "engineering", "reasoning_puzzle",
    "creative_writing", "professional_writing", "editing_grammar", "translation",
    "factual_qa", "education", "business", "finance", "legal", "health",
    "travel", "food", "entertainment", "lifestyle", "other",
)


@dataclass(frozen=True)
class RouterConfig:
    vocab_size: int = 4096
    dim: int = 128
    n_heads: int = 8
    layers: int = 3
    ffn_hidden: int = 208
    max_len: int = 256
    rope_base: float = 10_000.0
    proj_dim: int = 64
    n_domains: int = len(DOMAINS)


def build_rope_cache(config: RouterConfig) -> tuple[Tensor, Tensor]:
    head_dim = config.dim // config.n_heads
    positions = torch.arange(config.max_len, dtype=torch.float32)
    inv_freq = 1.0 / (config.rope_base ** (torch.arange(0, head_dim, 2).float() / 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)


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 RouterAttention(nn.Module):
    def __init__(self, config: RouterConfig) -> None:
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.dim // config.n_heads
        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)
        self.out = nn.Linear(config.dim, config.dim, bias=False)

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, attn_mask: Tensor) -> Tensor:
        batch, seq_len, dim = x.shape
        q, k, v = self.qkv(x).split(dim, dim=-1)
        shape = (batch, seq_len, self.n_heads, self.head_dim)
        q = apply_rope(q.view(shape).transpose(1, 2), cos, sin)
        k = apply_rope(k.view(shape).transpose(1, 2), cos, sin)
        v = v.view(shape).transpose(1, 2)
        attended = functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        return self.out(attended.transpose(1, 2).reshape(batch, seq_len, dim))


class RouterBlock(nn.Module):
    def __init__(self, config: RouterConfig) -> None:
        super().__init__()
        self.attn_norm = nn.RMSNorm(config.dim)
        self.attn = RouterAttention(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) -> Tensor:
        x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask)
        return x + self.ffn(self.ffn_norm(x))


class GlintRouter(nn.Module):
    def __init__(self, config: RouterConfig | None = None) -> None:
        super().__init__()
        config = config or RouterConfig()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.dim, padding_idx=PAD_ID)
        self.blocks = nn.ModuleList(RouterBlock(config) for _ in range(config.layers))
        self.final_norm = nn.RMSNorm(config.dim)
        cos, sin = build_rope_cache(config)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        pooled = 2 * config.dim
        self.domain_head = nn.Linear(pooled, config.n_domains)
        self.complexity_dir = nn.Linear(pooled, 1, bias=False)
        self.complexity_bias = nn.Parameter(torch.zeros(COMPLEXITY_LEVELS - 1))
        self.binary_head = nn.Linear(pooled, len(BINARY_FIELDS))
        self.route_head = nn.Linear(pooled, 1)
        self.proj_head = nn.Linear(pooled, config.proj_dim)
        self.register_buffer("temperature", torch.ones(3), persistent=True)

    def encode(self, tokens: Tensor) -> Tensor:
        valid = tokens != PAD_ID
        seq_len = tokens.shape[1]
        cos = self.rope_cos[:seq_len].to(self.embed.weight.dtype)
        sin = self.rope_sin[:seq_len].to(self.embed.weight.dtype)
        attn_mask = torch.zeros(tokens.shape, dtype=self.embed.weight.dtype,
                                device=tokens.device)
        attn_mask = attn_mask.masked_fill(~valid, float("-inf"))[:, None, None, :]
        x = self.embed(tokens)
        for block in self.blocks:
            x = block(x, cos, sin, attn_mask)
        x = self.final_norm(x)
        mask = valid.unsqueeze(-1)
        mean = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        maximum = torch.nan_to_num(x.masked_fill(~mask, float("-inf")).max(dim=1).values,
                                   neginf=0.0)
        return torch.cat((mean, maximum), dim=-1)

    def forward(self, tokens: Tensor) -> dict[str, Tensor]:
        pooled = self.encode(tokens)
        return {
            "domain": self.domain_head(pooled),
            "complexity": self.complexity_dir(pooled) + self.complexity_bias,
            "binary": self.binary_head(pooled),
            "route": self.route_head(pooled).squeeze(-1),
            "proj": functional.normalize(self.proj_head(pooled), dim=-1),
        }

    def calibrated(self, tokens: Tensor) -> dict[str, Tensor]:
        """probabilities with the fitted temperatures applied. the routing
        arithmetic in policy.py runs on these numbers, so they have to mean
        something. temperatures were fitted on held-out data after training."""
        out = self.forward(tokens)
        route_t, complexity_t, binary_t = self.temperature.unbind()
        return {
            "domain": out["domain"].softmax(dim=-1),
            "complexity": (out["complexity"] / complexity_t).sigmoid(),
            "binary": (out["binary"] / binary_t).sigmoid(),
            "route": (out["route"] / route_t).sigmoid(),
            "proj": out["proj"],
        }


def complexity_from_cumulative(probabilities: Tensor) -> Tensor:
    """coral decode. level = 1 + how many thresholds the prompt clears."""
    return 1 + (probabilities > 0.5).sum(dim=-1)


def encode_batch(tokenizer, texts: list[str], max_len: int) -> Tensor:
    """bos-prefixed, right-padded. long prompts lose their tail, because the
    instruction verb lives at the front and the pasted context lives at the back."""
    out = torch.full((len(texts), max_len), PAD_ID, dtype=torch.long)
    for row, text in enumerate(texts):
        ids = [BOS_ID] + tokenizer.encode(text).ids[: max_len - 1]
        out[row, : len(ids)] = torch.tensor(ids, dtype=torch.long)
    return out


def load_router(directory: Path = HERE, device: str = "cpu"):
    """returns (model, tokenizer). reads config.json + model.safetensors + tokenizer.json."""
    from safetensors.torch import load_file
    from tokenizers import Tokenizer

    directory = Path(directory)
    config = RouterConfig(**json.loads((directory / "config.json").read_text())["model"])
    model = GlintRouter(config).to(device)
    model.load_state_dict(load_file(directory / "model.safetensors", device=device))
    model.eval()
    return model, Tokenizer.from_file(str(directory / "tokenizer.json"))


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