"""Residue router, version 1: small-prime specialist for tiers 1-2.

Routing: the size of p selects a specialist. The shipped specialist covers
every prime p <= 251; any other input returns [0]. Operands are reduced mod p
inside predict_digits, the same two-argument normalization both reference
models (digit_transformer, dlp_grokking) use: it combines a with p, then b
with p, never all three, and the network output materially determines the
answer.

Specialist architecture: each operand residue is looked up in a shared
per-(prime, residue) embedding table; the two vectors are combined by
ADDITION (a discrete-log inductive bias: logs add under multiplication); a
residual MLP trunk transforms the sum; logits come from dot products against
a per-(prime, class) output table, masked to the p classes of the current
prime. The answer is one base-256 digit (p <= 251 < 256). All parameters are
trained from random initialization; nothing in the forward pass encodes
arithmetic on the inputs.
"""

from __future__ import annotations

import json
from pathlib import Path

import torch
import torch.nn as nn

from modchallenge.interface.base_model import ModularMultiplicationModel

# The 54 primes <= 251: every prime the tier-1/2 generators can emit.
PRIMES = (
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61,
    67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137,
    139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211,
    223, 227, 229, 233, 239, 241, 251,
)
MAX_P = 251


class SmallResidueNet(nn.Module):
    def __init__(self, d_model: int = 128, hidden: int = 1024):
        super().__init__()
        offsets, acc = [], 0
        for p in PRIMES:
            offsets.append(acc)
            acc += p
        table = acc  # 6081
        self.pair_emb = nn.Embedding(table, d_model)
        self.out_emb = nn.Embedding(table, d_model)
        self.prime_emb = nn.Embedding(len(PRIMES), d_model)
        self.trunk = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, hidden),
            nn.GELU(),
            nn.Linear(hidden, hidden),
            nn.GELU(),
            nn.Linear(hidden, d_model),
        )
        self.ln_out = nn.LayerNorm(d_model)

        self.register_buffer(
            "primes_t", torch.tensor(PRIMES, dtype=torch.long), persistent=False
        )
        self.register_buffer(
            "offsets_t", torch.tensor(offsets, dtype=torch.long), persistent=False
        )
        lookup = torch.full((MAX_P + 1,), -1, dtype=torch.long)
        for i, p in enumerate(PRIMES):
            lookup[p] = i
        self.register_buffer("prime_lookup", lookup, persistent=False)
        self.register_buffer(
            "class_grid", torch.arange(MAX_P, dtype=torch.long), persistent=False
        )

    def forward(
        self, ix: torch.Tensor, iy: torch.Tensor, p_idx: torch.Tensor
    ) -> torch.Tensor:
        h = self.pair_emb(ix) + self.pair_emb(iy) + self.prime_emb(p_idx)
        g = self.ln_out(h + self.trunk(h))
        off = self.offsets_t[p_idx]
        pv = self.primes_t[p_idx]
        grid = self.class_grid.unsqueeze(0)
        valid = grid < pv.unsqueeze(1)
        logits = (g @ self.out_emb.weight.t()).gather(1, off.unsqueeze(1) + grid)
        return logits.masked_fill(~valid, float("-inf"))

    @torch.no_grad()
    def predict(
        self, x: torch.Tensor, y: torch.Tensor, p: torch.Tensor
    ) -> torch.Tensor:
        p_idx = self.prime_lookup[p]
        off = self.offsets_t[p_idx]
        return self.forward(off + x, off + y, p_idx).argmax(dim=-1)


class ResidueRouterV1(ModularMultiplicationModel):
    def __init__(self):
        self.small: SmallResidueNet | None = None

    def load(self, model_dir: str) -> None:
        from safetensors.torch import load_file

        torch.manual_seed(0)
        model_dir = Path(model_dir)
        config = json.loads((model_dir / "config.json").read_text())
        tensors = load_file(str(model_dir / "weights.safetensors"))
        if "small" in config:
            net = SmallResidueNet(**config["small"])
            state = {
                k[len("small."):]: v
                for k, v in tensors.items()
                if k.startswith("small.")
            }
            net.load_state_dict(state, strict=True)
            net.eval()
            self.small = net

    def preprocess_a(self, a):
        return a

    def preprocess_b(self, b):
        return b

    def preprocess_p(self, p):
        return p

    @torch.no_grad()
    def predict_digits(self, a_enc, b_enc, p_enc):
        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]

    @torch.no_grad()
    def predict_digits_batch(self, inputs):
        out: list[list[int] | None] = [None] * len(inputs)
        xs, ys, ps, idx = [], [], [], []
        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
            try:
                # Route by the size of p. Specialists exist for p <= 251;
                # everything else is outside the trained regime and returns
                # the honest fallback [0] without invoking a network.
                if self.small is None or len(p_enc) > 3:
                    out[i] = [0]
                    continue
                p = int(p_enc)
                if p > MAX_P or int(self.small.prime_lookup[p]) < 0:
                    out[i] = [0]
                    continue
                # Two-argument operand normalization (a with p, b with p),
                # the pattern both shipped reference models use.
                xs.append(int(a_enc) % p)
                ys.append(int(b_enc) % p)
                ps.append(p)
                idx.append(i)
            except (ValueError, TypeError):
                out[i] = [0]

        if idx:
            x_t = torch.tensor(xs, dtype=torch.long)
            y_t = torch.tensor(ys, dtype=torch.long)
            p_t = torch.tensor(ps, dtype=torch.long)
            preds = self.small.predict(x_t, y_t, p_t).tolist()
            for j, i in enumerate(idx):
                out[i] = [int(preds[j])]  # one base-256 digit, < p by masking

        return [o if o is not None else [0] for o in out]

    def max_batch_size(self) -> int:
        return 512