"""Bit-serial learned reducer for the Modular Arithmetic Challenge.

A single shared, p-conditioned transition cell, applied in a fixed bit loop,
computes ``(a * b) mod p``. The cell learns the per-step transition

    s' = (2*s + d*x) mod p

(state ``s``, multiplicand ``x``, control bit ``d``, modulus ``p``). The Python
loop only sequences the bits most-significant-first (Horner form) -- the
explicitly-allowed recurrent / looped structure. No modular product is computed
in Python or in hand-coded tensor arithmetic: every reduction and the multiply
itself are produced by the trained cell. Randomising the weights collapses
accuracy to chance, which is the operational provenance test.

Pipeline per problem (a, b, p):

    reduce(a) = scan bits of a MSB-first with x=1               -> a mod p
    reduce(b) = scan bits of b MSB-first with x=1               -> b mod p
    multiply  = scan bits of (b mod p) MSB-first with x=(a mod p) -> (a*b) mod p

State is carried as bits between steps (no integer reconstruction inside the
loop); the harness decoder reconstructs the integer answer from the emitted
base-2 digits.

Regime: the cell is trained for primes ``p < 2^32`` and operands up to 96 bits
(tiers 1-4). Outside that regime the model abstains and emits ``[0]`` -- the
honest fallback -- rather than running the cell out of distribution.
"""

from __future__ import annotations

from pathlib import Path

import torch
from torch import nn

from modchallenge.interface.base_model import ModularMultiplicationModel

# State / modulus bit-width. Covers tiers 1-4: every prime there is < 2^32, and
# every residue is < p, so 32 bits hold both the state and the modulus features.
L = 32
# Tiers 1-4 operands are at most 96 bits. Beyond this we are out of regime.
MAX_OP_BITS = 96


def to_bits(vals: torch.Tensor, width: int = L) -> torch.Tensor:
    """Small non-negative ints -> (N, width) bit tensor, MSB-first.

    Used only on the modulus and the constant multiplicand x=1 (both small);
    this is representation, not arithmetic on the operands.
    """
    shifts = torch.arange(width - 1, -1, -1, device=vals.device)
    return (vals[:, None] >> shifts[None, :]) & 1


class Cell(nn.Module):
    """Shared per-step transition: (s_bits, x_bits, p_bits, d) -> next s_bits.

    The three bit-channels are read as a length-L sequence by a bidirectional
    GRU; the control bit d is injected as an embedding. The head emits one
    logit per output bit position. The same weights are used for the reduce
    steps (x=1) and the multiply steps (x = a mod p)."""

    def __init__(self, dmodel: int = 96, hidden: int = 128):
        super().__init__()
        self.in_proj = nn.Linear(3, dmodel)
        self.d_emb = nn.Embedding(2, dmodel)
        self.gru = nn.GRU(
            dmodel, hidden, num_layers=2, batch_first=True, bidirectional=True
        )
        self.head = nn.Linear(2 * hidden, 1)

    def forward(self, feat: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
        x = self.in_proj(feat) + self.d_emb(d)[:, None, :]
        h, _ = self.gru(x)
        return self.head(h).squeeze(-1)


def _bits_of(n: int) -> list[int]:
    """Non-negative int -> MSB-first bit list. Per-argument tokenisation."""
    if n <= 0:
        return [0]
    out: list[int] = []
    while n > 0:
        out.append(n & 1)
        n >>= 1
    out.reverse()
    return out


class BitSerialReducer(ModularMultiplicationModel):
    def __init__(self) -> None:
        self.model: Cell | None = None
        self.device: torch.device | None = None

    # -- lifecycle ------------------------------------------------------

    def load(self, model_dir: str) -> None:
        if torch.backends.mps.is_available():
            self.device = torch.device("mps")
        elif torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        ckpt = torch.load(
            Path(model_dir) / "weights.pt",
            map_location=self.device,
            weights_only=True,
        )
        self.model = Cell(**ckpt.get("config", {}))
        self.model.load_state_dict(ckpt["state_dict"])
        self.model.to(self.device)
        self.model.eval()

    # -- per-argument tokenisation (each hook sees only its own argument) --

    def preprocess_a(self, a):
        return _bits_of(int(a))

    def preprocess_b(self, b):
        return _bits_of(int(b))

    def preprocess_p(self, p):
        return int(p)

    # -- inference ------------------------------------------------------

    @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]] = [[0] for _ in inputs]
        idx, a_lists, b_lists, p_vals = [], [], [], []
        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
            p = int(p_enc)
            a_bits = list(a_enc)
            b_bits = list(b_enc)
            # Out of the trained regime (tiers 4+, or tier-0 giant operands):
            # abstain instead of running the cell out of distribution.
            if (
                p < 2
                or p >= (1 << L)
                or len(a_bits) > MAX_OP_BITS
                or len(b_bits) > MAX_OP_BITS
            ):
                continue
            idx.append(i)
            a_lists.append(a_bits)
            b_lists.append(b_bits)
            p_vals.append(p)

        if not idx:
            return out

        dev = self.device
        p_bits = to_bits(torch.tensor(p_vals, device=dev)).float()
        ra = self._reduce(a_lists, p_bits, dev)  # (N, L) residue bits
        rb = self._reduce(b_lists, p_bits, dev)
        prod = self._mul(ra, rb, p_bits)         # (N, L) answer bits

        prod_list = prod.long().tolist()
        for j, i in enumerate(idx):
            out[i] = [int(x) for x in prod_list[j]]
        return out

    def max_batch_size(self) -> int:
        return 256

    # -- internals ------------------------------------------------------

    def _step(self, s_bits, x_bits, p_bits, d):
        feat = torch.stack([s_bits, x_bits, p_bits], dim=-1)
        logits = self.model(feat, d)
        return (torch.sigmoid(logits) > 0.5).float()

    def _reduce(self, bit_lists, p_bits, dev):
        """Free-running Horner reduction of each operand to its residue mod p.

        x is fixed to 1; the control bit at each step is the operand bit.
        Leading zeros (from padding shorter operands to the batch width) keep
        the state at 0, so padding is harmless."""
        n = len(bit_lists)
        width = max(len(b) for b in bit_lists)
        padded = torch.zeros((n, width), dtype=torch.long, device=dev)
        for r, bl in enumerate(bit_lists):
            if bl:
                padded[r, width - len(bl):] = torch.tensor(
                    bl, dtype=torch.long, device=dev
                )
        s_bits = torch.zeros((n, L), device=dev)
        x_bits = to_bits(torch.ones(n, dtype=torch.long, device=dev)).float()
        for pos in range(width):  # MSB-first
            s_bits = self._step(s_bits, x_bits, p_bits, padded[:, pos])
        return s_bits

    def _mul(self, ra_bits, rb_bits, p_bits):
        """(a mod p) * (b mod p) mod p via Horner over the residue's L bits.

        x is fixed to (a mod p); the control bit at each step is a bit of
        (b mod p), scanned MSB-first."""
        n = ra_bits.shape[0]
        s_bits = torch.zeros((n, L), device=ra_bits.device)
        rb_long = rb_bits.long()
        for k in range(L):  # MSB-first over the 16-bit residue
            s_bits = self._step(s_bits, ra_bits, p_bits, rb_long[:, k])
        return s_bits