"""Scan-Register Machine submission for the Modular Arithmetic Challenge.

Two learned scan units (a carry-monoid adder and a comparison +
borrow conditional-subtract) composed by a bit-streaming double-and-add
schedule. Register width tracks limbs(p) + 1; operands are streamed as
bits, so operand size costs iterations, never width. At load time, the finite
learned cell domains are materialized into class transition / resolver tables
by running the shipped weights on every codebook cell and base-32 limb pair.
Inference then scans over those learned class IDs, and states are
argmax-discretized 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 macro-steps ----------------------------------------------

    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]

    # -- 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_has_bit = [any(row[i] for row in a_rows) for i in range(a_len)]
        b_has_bit = [any(row[i] for row in b_rows) for i in range(b_len)]
        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)
        started = False
        for i in range(a_len):  # A = a mod p
            column = a_bits[:, i]
            has_bit = a_has_bit[i]
            if not started and not has_bit:
                continue  # leading zero padding with r == 0 is inert
            started = True
            u = self._add(r, r)
            if has_bit:
                u = torch.where(column.unsqueeze(1), self._add(u, one), u)
            r = self._reduce(u, p_reg)
        addend = r

        r = torch.zeros(batch, width, dtype=torch.long, device=self.device)
        started = False
        for i in range(b_len):  # r = (A * b) mod p
            column = b_bits[:, i]
            has_bit = b_has_bit[i]
            if not started and not has_bit:
                continue
            started = True
            r = self._reduce(self._add(r, r), p_reg)
            if has_bit:
                r_add = self._reduce(self._add(r, addend), p_reg)
                r = torch.where(column.unsqueeze(1), r_add, r)

        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