"""Residue router, version 2: small-prime specialist (tiers 1-2) plus a lifted
local-step pipeline for tier 3.

Routing by the size of p:

  * p <= 251  (tiers 1-2): the v1 residue specialist. Each operand residue is
    looked up in a shared per-(prime, residue) 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
    a per-(prime, class) output table masked to the p classes of the current
    prime. The answer is a single base-256 digit (p <= 251 < 256).

  * 251 < p < 65536 (tier 3): two trained shared LOCAL-RULE step nets composed
    through fixed wiring. After the operands are reduced mod p (the same
    two-argument normalization both reference models use), x, y are 16-bit
    residues. A MULTIPLY step (the shared carry rule c' = floor((S+c)/2) over
    the carry-save column sums, composed closed-loop through a fixed parity
    readout) emits the exact 32-bit product t = x*y. A REDUCTION step (a shared
    per-nibble borrow/compare rule, composed through fixed restoring-division
    wiring) emits r = t mod p. The answer is r, emitted as base-256 digits
    MSB-first (two digits cover a 16-bit residue). Both step nets are trained
    from random init; randomizing either one's weights collapses the pipeline.

  * p >= 65536 (tiers 4-10): outside the trained regime; returns [0].

Nothing in the forward pass hand-codes the arithmetic over the actual (a, b, p):
the carry-save column sums, the parity readout, the bit shifts, the restoring-
division topology, and the ge-from-final-borrow decision are FIXED scaffold; the
two NONTRIVIAL decisions -- the carry rule and the borrow/compare rule -- live
in trained MLP parameters. The output digits materially determine the answer.
"""

from __future__ import annotations

import json
from pathlib import Path

import torch
import torch.nn as nn

from modchallenge.interface.base_model import ModularMultiplicationModel

# ===========================================================================
# Tier 1-2 specialist (v1 residue net), vendored verbatim
# ===========================================================================

# 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)


# ===========================================================================
# Tier 3 step nets + fixed wiring (vendored from t3_step_model)
# ===========================================================================

# -- multiply step geometry (16x16 -> 32-bit) -------------------------------
MUL_N_OPERAND_BITS = 16
MUL_N_PRODUCT_BITS = 32
MUL_N_COLUMNS = 2 * MUL_N_OPERAND_BITS - 1        # 31 partial-product columns
MUL_N_CARRIES = MUL_N_PRODUCT_BITS - 1            # carries into columns 1..31
MUL_SUM_BITS = 5                                  # S_c <= 16 needs 5 bits
MUL_CARRY_BITS = 4                                # carry over the chain <= 15
MUL_STEP_IN = MUL_SUM_BITS + MUL_CARRY_BITS       # 9

# -- reduction step geometry (nibble borrow ripple) -------------------------
NIB = 4                                            # nibble width in bits
RED_NIBBLES = 5                                    # 17-bit R_pre / 16-bit p -> 5 nibbles
RED_STEP_IN = NIB + NIB + 1                        # a_nib(4) + b_nib(4) + borrow_in(1)
RED_STEP_OUT = NIB + 1                             # diff nibble(4) + borrow_out(1)
T_BITS = 32                                         # t = x * y is 32-bit


class MulCarryStep(nn.Module):
    """Shared carry step for 16-bit multiply: 9-bit local state -> 4 carry bits."""

    def __init__(self, width: int = 96, depth: int = 3):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(MUL_STEP_IN, width)])
        for _ in range(depth - 1):
            self.layers.append(nn.Linear(width, width))
        self.head = nn.Linear(width, MUL_CARRY_BITS)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = x
        for lin in self.layers:
            h = self.act(lin(h))
        return self.head(h)


class RedBorrowStep(nn.Module):
    """Shared reduction nibble step: 9-bit local state -> 5 bits (diff+borrow)."""

    def __init__(self, width: int = 96, depth: int = 3):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(RED_STEP_IN, width)])
        for _ in range(depth - 1):
            self.layers.append(nn.Linear(width, width))
        self.head = nn.Linear(width, RED_STEP_OUT)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = x
        for lin in self.layers:
            h = self.act(lin(h))
        return self.head(h)


def _bits16(v: torch.Tensor) -> torch.Tensor:
    return ((v.unsqueeze(1) >> torch.arange(16, device=v.device)) & 1).float()


def _column_sums_16(x_bits: torch.Tensor, y_bits: torch.Tensor) -> torch.Tensor:
    """(N,16),(N,16) operand bits -> (N,31) carry-save column sums (FIXED scaffold)."""
    outer = x_bits.unsqueeze(2) * y_bits.unsqueeze(1)        # (N,16,16)
    n = outer.shape[0]
    s = torch.zeros(n, MUL_N_COLUMNS, dtype=outer.dtype, device=outer.device)
    for i in range(MUL_N_OPERAND_BITS):
        for j in range(MUL_N_OPERAND_BITS):
            s[:, i + j] += outer[:, i, j]
    return s


def _encode_carry_inputs(s: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
    """(N,) sums and carries -> (N, 9) float bits, LSB first."""
    si = torch.arange(MUL_SUM_BITS, device=s.device)
    ci = torch.arange(MUL_CARRY_BITS, device=c.device)
    sb = ((s.unsqueeze(1) >> si) & 1).float()
    cb = ((c.unsqueeze(1) >> ci) & 1).float()
    return torch.cat([sb, cb], dim=1)


def _carry_bits_to_int(bits: torch.Tensor) -> torch.Tensor:
    w = (1 << torch.arange(MUL_CARRY_BITS, device=bits.device)).long()
    return (bits.round().clamp(0, 1).long() * w).sum(dim=-1)


def _routed_product_logits(carry_logits45: torch.Tensor, col_parity: torch.Tensor) -> torch.Tensor:
    """Fixed parity readout (no parameters): carry-bit logits + parity -> 32 bit-logits."""
    BIG = 20.0
    lsb = carry_logits45[:, 0::MUL_CARRY_BITS]              # (B, 31) lsb of c_1..c_31
    bit0 = (2.0 * col_parity[:, 0:1] - 1.0) * BIG
    mid = (1.0 - 2.0 * col_parity[:, 1:]) * lsb[:, :-1]     # bits 1..30
    bit31 = lsb[:, -1:]
    return torch.cat([bit0, mid, bit31], dim=1)


@torch.no_grad()
def _closed_loop_mul(step: nn.Module, col_sums: torch.Tensor) -> torch.Tensor:
    """Compose the trained carry step over 31 columns -> carry-bit logits (B, 31*4)."""
    n = col_sums.shape[0]
    s = col_sums.long()
    carry = torch.zeros(n, dtype=torch.long, device=s.device)
    out = torch.empty(n, MUL_N_CARRIES * MUL_CARRY_BITS, device=col_sums.device)
    for c in range(MUL_N_COLUMNS):
        lg = step(_encode_carry_inputs(s[:, c], carry))
        out[:, MUL_CARRY_BITS * c:MUL_CARRY_BITS * (c + 1)] = lg
        carry = _carry_bits_to_int((lg > 0).float())
    return out


@torch.no_grad()
def _composed_product(step: nn.Module, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Trained carry step (closed loop) + fixed parity readout -> 32-bit product (B,)."""
    col_sums = _column_sums_16(_bits16(x), _bits16(y))
    logits = _closed_loop_mul(step, col_sums)
    col_parity = (col_sums.long() & 1).float()
    bit_logits = _routed_product_logits(logits, col_parity)
    bits = (bit_logits > 0).long()
    w = (1 << torch.arange(MUL_N_PRODUCT_BITS, device=bits.device)).long()
    return (bits * w).sum(dim=1)


def _encode_red_inputs(a: torch.Tensor, b: torch.Tensor, bin_: torch.Tensor) -> torch.Tensor:
    """(N,) nibbles + borrow -> (N, 9) float bits (a nib LSB first, b nib, borrow)."""
    ai = torch.arange(NIB, device=a.device)
    aa = ((a.unsqueeze(1) >> ai) & 1).float()
    bb = ((b.unsqueeze(1) >> ai) & 1).float()
    cc = bin_.float().unsqueeze(1)
    return torch.cat([aa, bb, cc], dim=1)


def _red_bits_to_out(bits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """(N,5) logits-thresholded bits -> (diff nibble int, borrow_out int)."""
    hb = (bits > 0).long()
    w = (1 << torch.arange(NIB, device=bits.device)).long()
    d = (hb[:, :NIB] * w).sum(dim=1)
    bout = hb[:, NIB]
    return d, bout


@torch.no_grad()
def _composed_reduce(step: nn.Module, t: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
    """Trained borrow step composed through fixed restoring-division wiring -> r (B,).

    The bit shifts, the ge-from-final-borrow decision, and the keep/replace of R
    are fixed scaffold; the per-nibble subtract DECISION is the trained step.
    """
    n = t.shape[0]
    device = t.device
    R = torch.zeros(n, dtype=torch.long, device=device)
    p_nib = torch.stack([(p >> (NIB * k)) & 0xF for k in range(RED_NIBBLES)], dim=1)
    wk = (1 << (NIB * torch.arange(RED_NIBBLES, device=device))).long()
    for i in range(T_BITS - 1, -1, -1):
        bit = (t >> i) & 1
        Rpre = (R << 1) | bit
        borrow = torch.zeros(n, dtype=torch.long, device=device)
        diff_nib = torch.zeros(n, RED_NIBBLES, dtype=torch.long, device=device)
        for k in range(RED_NIBBLES):
            an = (Rpre >> (NIB * k)) & 0xF
            bn = p_nib[:, k]
            lg = step(_encode_red_inputs(an, bn, borrow))
            d, bout = _red_bits_to_out(lg)
            diff_nib[:, k] = d
            borrow = bout
        ge = (borrow == 0).long()
        diff_val = (diff_nib * wk).sum(dim=1)
        R = torch.where(ge.bool(), diff_val, Rpre)
    return R


# ===========================================================================
# Router
# ===========================================================================

T3_MIN_P = MAX_P + 1     # 252: first prime size routed to the lifted pipeline
T3_MAX_P = (1 << 16) - 1  # tier-3 primes are 9-16 bits


class ResidueRouterV1(ModularMultiplicationModel):
    """Router over per-tier specialists, selected by the size of p.

    Kept the class name ``ResidueRouterV1`` so the manifest entry_class is
    stable across versions; this is v2 (tiers 1-3).
    """

    def __init__(self):
        self.small: SmallResidueNet | None = None
        self.mul: MulCarryStep | None = None
        self.red: RedBorrowStep | 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())

        # tier 1-2 specialist
        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

        # tier 3 lifted step nets
        if "t3" in config:
            arch = config["t3"]
            mul = MulCarryStep(width=arch["width"], depth=arch["depth"])
            red = RedBorrowStep(width=arch["width"], depth=arch["depth"])
            mul.load_state_dict(load_file(str(model_dir / "t3_mul.safetensors")), strict=True)
            red.load_state_dict(load_file(str(model_dir / "t3_red.safetensors")), strict=True)
            mul.eval()
            red.eval()
            self.mul = mul
            self.red = red

    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)
        # tier 1-2 batch (single base-256 digit)
        s_x, s_y, s_p, s_idx = [], [], [], []
        # tier 3 batch (two base-256 digits)
        t_x, t_y, t_p, t_idx = [], [], [], []

        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
            try:
                p = int(p_enc)
            except (ValueError, TypeError):
                out[i] = [0]
                continue
            # Operand normalization: combine a with p, then b with p (the
            # two-argument reduction both reference models use). Never all three.
            try:
                xr = int(a_enc) % p
                yr = int(b_enc) % p
            except (ValueError, TypeError):
                out[i] = [0]
                continue

            if self.small is not None and 2 <= p <= MAX_P and int(self.small.prime_lookup[p]) >= 0:
                s_x.append(xr); s_y.append(yr); s_p.append(p); s_idx.append(i)
            elif self.mul is not None and T3_MIN_P <= p <= T3_MAX_P:
                t_x.append(xr); t_y.append(yr); t_p.append(p); t_idx.append(i)
            else:
                # outside the trained regime (tiers 4-10) -> honest fallback
                out[i] = [0]

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

        if t_idx:
            x_t = torch.tensor(t_x, dtype=torch.long)
            y_t = torch.tensor(t_y, dtype=torch.long)
            p_t = torch.tensor(t_p, dtype=torch.long)
            prod = _composed_product(self.mul, x_t, y_t)       # exact 32-bit t
            r = _composed_reduce(self.red, prod, p_t)          # r = t mod p in [0, p)
            r_list = r.tolist()
            for j, i in enumerate(t_idx):
                rv = int(r_list[j])
                # base-256 digits, MSB-first (two digits cover a 16-bit residue)
                out[i] = [rv >> 8, rv & 0xFF]

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

    def max_batch_size(self) -> int:
        return 512