"""Uniform-transition Scan-Register Machine for the Modular Arithmetic Challenge.

Every raw bit of both operands is processed by the same learned transition
over base-32 registers. The transition is built from a learned carry-monoid
adder and a learned comparison/borrow conditional subtractor. Register width
tracks limbs(p) + 1; operand size costs iterations, never width. At load time,
the complete finite learned cell domains are materialized into class-transition
and resolver tables by evaluating the shipped weights. Inference scans over
those learned class IDs and argmax-discretizes the register between steps.

No big-integer arithmetic computes the answer at inference: Python ints are
used only inside the per-argument preprocess hooks to convert each decimal
string into base-32 limbs / bits (base conversion, explicitly allowed).
The emitted digits come from the learned resolvers; randomizing the weights
collapses accuracy to chance.
"""

from __future__ import annotations

from pathlib import Path

import torch
from torch import nn

from modchallenge.interface.base_model import ModularMultiplicationModel

BASE = 32
BITS_PER_LIMB = 5


def _mlp(d_in: int, hidden: int, d_out: int) -> nn.Sequential:
    return nn.Sequential(nn.Linear(d_in, hidden), nn.GELU(), nn.Linear(hidden, d_out))


def scan_tree(compose, identity, sig):
    batch, n, d = sig.shape
    size = 1
    while size < n:
        size *= 2
    buf = torch.empty(batch, size, d, device=sig.device, dtype=sig.dtype)
    buf[:, :n] = sig
    if size > n:
        buf[:, n:] = identity
    stride = 1
    while stride < size:
        buf[:, 2 * stride - 1 :: 2 * stride] = compose(
            buf[:, stride - 1 :: 2 * stride], buf[:, 2 * stride - 1 :: 2 * stride]
        )
        stride *= 2
    total = buf[:, -1].clone()
    buf[:, -1] = identity
    stride = size // 2
    while stride >= 1:
        left = buf[:, stride - 1 :: 2 * stride].clone()
        parent = buf[:, 2 * stride - 1 :: 2 * stride].clone()
        buf[:, stride - 1 :: 2 * stride] = parent
        buf[:, 2 * stride - 1 :: 2 * stride] = compose(parent, left)
        stride //= 2
    return buf[:, :n], total


def scan_tree_classes(op_table: torch.Tensor, identity: int, sig: torch.Tensor):
    """Blelloch exclusive scan over finite learned cell IDs."""
    batch, n = sig.shape
    size = 1
    while size < n:
        size *= 2
    buf = torch.full((batch, size), identity, device=sig.device, dtype=torch.long)
    buf[:, :n] = sig
    stride = 1
    while stride < size:
        left = buf[:, stride - 1 :: 2 * stride]
        right = buf[:, 2 * stride - 1 :: 2 * stride]
        buf[:, 2 * stride - 1 :: 2 * stride] = op_table[left, right]
        stride *= 2
    total = buf[:, -1].clone()
    buf[:, -1] = identity
    stride = size // 2
    while stride >= 1:
        left = buf[:, stride - 1 :: 2 * stride].clone()
        parent = buf[:, 2 * stride - 1 :: 2 * stride].clone()
        buf[:, stride - 1 :: 2 * stride] = parent
        buf[:, 2 * stride - 1 :: 2 * stride] = op_table[parent, left]
        stride //= 2
    return buf[:, :n], total


def reduce_total(compose, identity, sig):
    """Root of the scan tree only (for the comparison verdict)."""
    x = sig
    batch, _, d = sig.shape
    while x.shape[1] > 1:
        if x.shape[1] % 2:
            x = torch.cat([x, identity.expand(batch, 1, d)], dim=1)
        x = compose(x[:, 0::2], x[:, 1::2])
    return x[:, 0]


class ScanAdder(nn.Module):
    def __init__(self, base: int = BASE, d_emb: int = 32, d_sig: int = 16, hidden: int = 96):
        super().__init__()
        self.limb_emb = nn.Embedding(base, d_emb)
        self.encoder = _mlp(2 * d_emb, hidden, d_sig)
        self.op = _mlp(2 * d_sig, hidden, d_sig)
        self.identity = nn.Parameter(torch.zeros(d_sig))
        self.resolver = _mlp(2 * d_emb + d_sig, hidden, base)
        self.carry_head = _mlp(d_sig, hidden, 2)

    def compose(self, left, right):
        return self.op(torch.cat([left, right], dim=-1))


class ModReduce(nn.Module):
    def __init__(self, base: int = BASE, d_emb: int = 32, d_sig: int = 16, hidden: int = 96):
        super().__init__()
        self.limb_emb = nn.Embedding(base, d_emb)
        self.cmp_encoder = _mlp(2 * d_emb, hidden, d_sig)
        self.cmp_op = _mlp(2 * d_sig, hidden, d_sig)
        self.cmp_identity = nn.Parameter(torch.zeros(d_sig))
        self.borrow_encoder = _mlp(2 * d_emb, hidden, d_sig)
        self.borrow_op = _mlp(2 * d_sig, hidden, d_sig)
        self.borrow_identity = nn.Parameter(torch.zeros(d_sig))
        self.resolver = _mlp(2 * d_emb + 2 * d_sig, hidden, base)
        self.sub_head = _mlp(d_sig, hidden, 2)
        self.borrow_head = _mlp(d_sig, hidden, 2)

    def compose_cmp(self, left, right):
        return self.cmp_op(torch.cat([left, right], dim=-1))

    def compose_borrow(self, left, right):
        return self.borrow_op(torch.cat([left, right], dim=-1))


def _snap(vec: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor:
    dist = torch.cdist(vec.reshape(-1, vec.shape[-1]), codebook)
    return codebook[dist.argmin(dim=-1)].reshape(vec.shape)


def _nearest_class(vec: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor:
    flat = vec.reshape(-1, vec.shape[-1])
    scores = (
        flat.square().sum(-1, keepdim=True)
        - 2 * flat @ codebook.T
        + codebook.square().sum(-1).unsqueeze(0)
    )
    return scores.argmin(dim=-1).reshape(vec.shape[:-1])


class ScanRegisterMachine(ModularMultiplicationModel):
    """Entry class declared in manifest.json."""

    def load(self, model_dir: str) -> None:
        directory = Path(model_dir)
        if torch.cuda.is_available():
            self.device = "cuda"
        elif torch.backends.mps.is_available():
            self.device = "mps"
        else:
            self.device = "cpu"
        self.adder = ScanAdder()
        self.adder.load_state_dict(
            torch.load(directory / "adder.pt", map_location="cpu")
        )
        self.reducer = ModReduce()
        self.reducer.load_state_dict(
            torch.load(directory / "reducer.pt", map_location="cpu")
        )
        self.adder.to(self.device).eval()
        self.reducer.to(self.device).eval()
        codebooks = torch.load(directory / "codebooks.pt", map_location="cpu")
        self.carry_cb = codebooks["carry"].to(self.device)
        self.cmp_cb = codebooks["cmp"].to(self.device)
        self.borrow_cb = codebooks["borrow"].to(self.device)
        self.carry_identity = self.carry_cb[1]  # propagate
        self.cmp_identity = self.cmp_cb[1]  # EQ
        self.borrow_identity = self.borrow_cb[1]  # propagate
        self._build_tables()
        torch.set_grad_enabled(False)

    def _build_tables(self) -> None:
        digits = torch.arange(BASE, device=self.device)

        adder_emb = self.adder.limb_emb(digits)
        x = digits.repeat_interleave(BASE)
        y = digits.repeat(BASE)
        adder_pair = torch.cat([adder_emb[x], adder_emb[y]], dim=-1)
        adder_sig = self.adder.encoder(adder_pair)
        self.adder_pair_class = _nearest_class(adder_sig, self.carry_cb).reshape(
            BASE, BASE
        )

        carry_count = self.carry_cb.shape[0]
        left = self.carry_cb[:, None, :].expand(carry_count, carry_count, -1)
        right = self.carry_cb[None, :, :].expand(carry_count, carry_count, -1)
        carry_out = self.adder.compose(
            left.reshape(carry_count * carry_count, -1),
            right.reshape(carry_count * carry_count, -1),
        )
        self.carry_op_table = _nearest_class(carry_out, self.carry_cb).reshape(
            carry_count, carry_count
        )

        adder_pair_exp = adder_pair[:, None, :].expand(BASE * BASE, carry_count, -1)
        carry_exp = self.carry_cb[None, :, :].expand(BASE * BASE, carry_count, -1)
        adder_logits = self.adder.resolver(
            torch.cat([adder_pair_exp, carry_exp], dim=-1).reshape(
                BASE * BASE * carry_count, -1
            )
        )
        self.adder_digit_table = adder_logits.argmax(-1).reshape(
            BASE, BASE, carry_count
        )

        reducer_emb = self.reducer.limb_emb(digits)
        reducer_pair = torch.cat([reducer_emb[x], reducer_emb[y]], dim=-1)
        cmp_sig = self.reducer.cmp_encoder(reducer_pair)
        self.cmp_pair_class = _nearest_class(cmp_sig, self.cmp_cb).reshape(BASE, BASE)

        cmp_count = self.cmp_cb.shape[0]
        left = self.cmp_cb[:, None, :].expand(cmp_count, cmp_count, -1)
        right = self.cmp_cb[None, :, :].expand(cmp_count, cmp_count, -1)
        cmp_out = self.reducer.compose_cmp(
            left.reshape(cmp_count * cmp_count, -1),
            right.reshape(cmp_count * cmp_count, -1),
        )
        self.cmp_op_table = _nearest_class(cmp_out, self.cmp_cb).reshape(
            cmp_count, cmp_count
        )

        borrow_count = self.borrow_cb.shape[0]
        borrow_sig = self.reducer.borrow_encoder(reducer_pair)
        self.borrow_pair_class = _nearest_class(
            borrow_sig, self.borrow_cb
        ).reshape(BASE, BASE)

        left = self.borrow_cb[:, None, :].expand(
            borrow_count, borrow_count, -1
        )
        right = self.borrow_cb[None, :, :].expand(
            borrow_count, borrow_count, -1
        )
        borrow_out = self.reducer.compose_borrow(
            left.reshape(borrow_count * borrow_count, -1),
            right.reshape(borrow_count * borrow_count, -1),
        )
        self.borrow_op_table = _nearest_class(
            borrow_out, self.borrow_cb
        ).reshape(borrow_count, borrow_count)

        reducer_pair_exp = reducer_pair[:, None, None, :].expand(
            BASE * BASE, borrow_count, cmp_count, -1
        )
        borrow_exp = self.borrow_cb[None, :, None, :].expand(
            BASE * BASE, borrow_count, cmp_count, -1
        )
        verdict_exp = self.cmp_cb[None, None, :, :].expand(
            BASE * BASE, borrow_count, cmp_count, -1
        )
        reducer_logits = self.reducer.resolver(
            torch.cat([reducer_pair_exp, borrow_exp, verdict_exp], dim=-1).reshape(
                BASE * BASE * borrow_count * cmp_count, -1
            )
        )
        self.reducer_digit_table = reducer_logits.argmax(-1).reshape(
            BASE, BASE, borrow_count, cmp_count
        )

    # -- per-argument preprocessing (base conversion only) ----------------

    def preprocess_a(self, a: str) -> list[int]:
        value = int(a)
        bits = []
        while value:
            bits.append(value & 1)
            value >>= 1
        return list(reversed(bits)) or [0]  # MSB-first

    def preprocess_b(self, b: str) -> list[int]:
        return self.preprocess_a(b)

    def preprocess_p(self, p: str) -> list[int]:
        value = int(p)
        limbs = []
        while value:
            limbs.append(value % BASE)
            value //= BASE
        return limbs or [0]  # LSB-first

    # -- learned finite primitives ---------------------------------------

    def _add(self, x, y):
        sig = self.adder_pair_class[x, y]
        prefixes, _ = scan_tree_classes(self.carry_op_table, 1, sig)
        return self.adder_digit_table[x, y, prefixes]

    def _reduce(self, u, p_reg):
        """Learned comparison and borrow scans for conditional subtraction."""
        cmp_sig = self.cmp_pair_class[u, p_reg]
        # Comparison is most-significant-first, so only the total over the
        # reversed limb stream is needed for the subtract/keep verdict.
        _, total = scan_tree_classes(
            self.cmp_op_table, 1, cmp_sig.flip(1)
        )

        # Borrow propagation is least-significant-first. Both its local
        # classifications and its composition table come from the trained
        # reducer weights; no arithmetic class mapping is supplied by Python.
        borrow_sig = self.borrow_pair_class[u, p_reg]
        borrow_index, _ = scan_tree_classes(
            self.borrow_op_table, 1, borrow_sig
        )
        verdict = total.unsqueeze(1).expand_as(u)
        return self.reducer_digit_table[u, p_reg, borrow_index, verdict]

    def _step(self, r, addend, p_reg, bit):
        """One uniform learned transition for every raw operand bit.

        The external loop does not choose an arithmetic routine from the bit
        or phase. It always applies this same transition. Two learned
        conditional-subtract passes cover the complete transition range
        ``2*r + bit*addend < 3*p``.
        """
        zero = torch.zeros_like(addend)
        token_addend = torch.where(bit.unsqueeze(1), addend, zero)
        u = self._add(self._add(r, r), token_addend)
        return self._reduce(self._reduce(u, p_reg), p_reg)

    # -- the machine --------------------------------------------------------

    @torch.inference_mode()
    def _rollout(self, group: list[tuple[list[int], list[int], list[int]]]) -> list[list[int]]:
        width = max(len(p) for _, _, p in group) + 1
        a_len = max(len(a) for a, _, _ in group)
        b_len = max(len(b) for _, b, _ in group)
        batch = len(group)
        pad_bits = lambda bits, n: [0] * (n - len(bits)) + bits
        a_rows = [pad_bits(a, a_len) for a, _, _ in group]
        b_rows = [pad_bits(b, b_len) for _, b, _ in group]
        a_bits = torch.tensor(
            a_rows, dtype=torch.bool, device=self.device
        )
        b_bits = torch.tensor(
            b_rows, dtype=torch.bool, device=self.device
        )
        p_reg = torch.tensor(
            [p + [0] * (width - len(p)) for _, _, p in group],
            dtype=torch.long,
            device=self.device,
        )
        one = torch.zeros(batch, width, dtype=torch.long, device=self.device)
        one[:, 0] = 1

        r = torch.zeros(batch, width, dtype=torch.long, device=self.device)
        for i in range(a_len):
            r = self._step(r, one, p_reg, a_bits[:, i])
        addend = r

        r = torch.zeros(batch, width, dtype=torch.long, device=self.device)
        for i in range(b_len):
            r = self._step(r, addend, p_reg, b_bits[:, i])

        rows = r.cpu().tolist()
        return [[int(d) for d in reversed(row)] for row in rows]  # MSB-first

    # -- interface -----------------------------------------------------------

    def predict_digits(self, a_enc, b_enc, p_enc) -> list[int]:
        self._build_tables()
        return self._rollout([(a_enc, b_enc, p_enc)])[0]

    def predict_digits_batch(self, inputs) -> list[list[int]]:
        self._build_tables()
        # group by size bucket so short problems don't pay for long ones
        groups: dict[tuple[int, int, int], list[int]] = {}
        for index, (a, b, p) in enumerate(inputs):
            key = (
                -(-len(p) // 8),
                -(-len(a) // 64),
                -(-len(b) // 64),
            )
            groups.setdefault(key, []).append(index)
        results: list[list[int] | None] = [None] * len(inputs)
        for indices in groups.values():
            rows = self._rollout([inputs[i] for i in indices])
            for i, row in zip(indices, rows):
                results[i] = row
        return results  # type: ignore[return-value]

    def max_batch_size(self) -> int:
        return 100