"""MiniNeuralHorner development model for modular multiplication.

The learned component is one modulus-conditioned recurrent transition cell.
It predicts the next binary residue state for

    s_next = (2 * s + d * x) mod p.

A fixed Horner schedule applies that cell to reduce both operands and then
multiply the two residues. The emitted answer is a list of base-2 digits. The
SAIR evaluation harness performs the final digit decoding.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import torch
from torch import nn

from modchallenge.interface.base_model import ModularMultiplicationModel

_MASK32 = (1 << 32) - 1
_CHECKPOINT_SCHEMA = "mini-neuralhorner-inference-v1"
_EXPECTED_CONFIG = {
    "bidirectional": True,
    "dmodel": 96,
    "hidden": 61,
    "num_layers": 2,
}
_EXPECTED_PARAMETERS = 126_603
_EXPECTED_WIDTH = 2_048
_EXPECTED_TENSOR_SHA256 = (
    "7d1768ae1260f750e0a80ec93d98f86a80d441e479ce29e0a8e21fd098c742a3"
)
_QUALIFIED_CHECKPOINT_SHA256 = (
    "d296b711bb6a7faaa1dd81e05478cfa75f11071c42a8c36fbf60e758ee7eb407"
)


def _to_bits_small(values: torch.Tensor, width: int) -> torch.Tensor:
    shifts = torch.arange(width - 1, -1, -1, device=values.device)
    return (values[:, None] >> shifts[None, :]) & 1


def to_bits_limbs(values: list[int], device: torch.device, width: int) -> torch.Tensor:
    """Convert nonnegative Python integers to MSB-first bits without int64 overflow."""
    limb_count = (width + 31) // 32
    columns = []
    for limb_index in range(limb_count - 1, -1, -1):
        limb = torch.tensor(
            [(value >> (32 * limb_index)) & _MASK32 for value in values],
            dtype=torch.int64,
            device=device,
        )
        columns.append(_to_bits_small(limb, 32))
    bits = torch.cat(columns, dim=1)
    excess = limb_count * 32 - width
    return bits[:, excess:] if excess else bits


def _tensor_digest(state_dict: dict[str, torch.Tensor]) -> str:
    """Hash tensor names, dtypes, shapes, and raw values in a stable order."""
    digest = hashlib.sha256()
    for name in sorted(state_dict):
        tensor = state_dict[name].detach().cpu().contiguous()
        header = json.dumps(
            {"dtype": str(tensor.dtype), "name": name, "shape": list(tensor.shape)},
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
        raw = tensor.numpy().tobytes(order="C")
        digest.update(len(header).to_bytes(8, "big"))
        digest.update(header)
        digest.update(len(raw).to_bytes(8, "big"))
        digest.update(raw)
    return digest.hexdigest()


class TransitionCell(nn.Module):
    def __init__(
        self,
        dmodel: int,
        hidden: int,
        num_layers: int,
        bidirectional: bool,
    ) -> None:
        super().__init__()
        directions = 2 if bidirectional else 1
        self.in_proj = nn.Linear(3, dmodel)
        self.d_emb = nn.Embedding(2, dmodel)
        self.gru = nn.GRU(
            dmodel,
            hidden,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=bidirectional,
        )
        self.head = nn.Linear(directions * hidden, 1)

    def forward(
        self,
        features: torch.Tensor,
        control: torch.Tensor,
    ) -> torch.Tensor:
        embedded = self.in_proj(features) + self.d_emb(control)[:, None, :]
        hidden, _ = self.gru(embedded)
        return self.head(hidden).squeeze(-1)


def _bits_of(value: int) -> list[int]:
    if value <= 0:
        return [0]
    digits = []
    while value:
        digits.append(value & 1)
        value >>= 1
    digits.reverse()
    return digits


class MiniNeuralHorner(ModularMultiplicationModel):
    """SAIR interface adapter for the 126,603-parameter transition cell."""

    def __init__(self) -> None:
        self.model: TransitionCell | None = None
        self.device = torch.device("cpu")
        self.width = _EXPECTED_WIDTH
        self._sequence_width = 32

    def load(self, model_dir: str) -> None:
        checkpoint_path = Path(model_dir) / "weights.pt"
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
        if checkpoint.get("schema") != _CHECKPOINT_SCHEMA:
            raise ValueError("unsupported MiniNeuralHorner checkpoint schema")
        if checkpoint.get("config") != _EXPECTED_CONFIG:
            raise ValueError("checkpoint architecture does not match the packaged model")
        if checkpoint.get("L") != _EXPECTED_WIDTH:
            raise ValueError("checkpoint inference width does not match the packaged model")
        provenance = checkpoint.get("provenance", {})
        if provenance.get("qualified_checkpoint_sha256") != _QUALIFIED_CHECKPOINT_SHA256:
            raise ValueError("checkpoint provenance does not match the qualified source")

        state_dict = checkpoint.get("state_dict")
        if not isinstance(state_dict, dict):
            raise ValueError("checkpoint is missing its state_dict")
        parameter_count = sum(tensor.numel() for tensor in state_dict.values())
        if parameter_count != _EXPECTED_PARAMETERS:
            raise ValueError("checkpoint parameter count is invalid")
        if _tensor_digest(state_dict) != _EXPECTED_TENSOR_SHA256:
            raise ValueError("checkpoint tensor digest is invalid")

        if torch.cuda.is_available():
            self.device = torch.device("cuda")
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.allow_tf32 = False
            torch.backends.cuda.matmul.allow_tf32 = False
        elif torch.backends.mps.is_available():
            self.device = torch.device("mps")

        self.model = TransitionCell(**_EXPECTED_CONFIG)
        self.model.load_state_dict(state_dict, strict=True)
        self.model.to(self.device)
        self.model.eval()

    def preprocess_a(self, a: str) -> list[int]:
        return _bits_of(int(a))

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

    def preprocess_p(self, p: str) -> int:
        return int(p)

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

    @torch.no_grad()
    def predict_digits_batch(
        self,
        inputs: list[tuple[list[int], list[int], int]],
    ) -> list[list[int]]:
        if self.model is None:
            raise RuntimeError("load() must be called before inference")

        max_operand_bits = 4 * self.width
        outputs: list[list[int]] = [[0] for _ in inputs]
        valid_indices = []
        a_bit_lists = []
        b_bit_lists = []
        moduli = []
        for index, (a_enc, b_enc, p_enc) in enumerate(inputs):
            modulus = int(p_enc)
            a_bits = list(a_enc)
            b_bits = list(b_enc)
            if (
                modulus < 2
                or modulus >= (1 << self.width)
                or len(a_bits) > max_operand_bits
                or len(b_bits) > max_operand_bits
            ):
                continue
            valid_indices.append(index)
            a_bit_lists.append(a_bits)
            b_bit_lists.append(b_bits)
            moduli.append(modulus)

        if not valid_indices:
            return outputs

        maximum_modulus_bits = max(modulus.bit_length() for modulus in moduli)
        self._sequence_width = min(
            self.width,
            max(32, ((maximum_modulus_bits + 31) // 32) * 32),
        )
        modulus_bits = to_bits_limbs(
            moduli,
            self.device,
            self._sequence_width,
        ).float()
        a_residues = self._reduce(a_bit_lists, modulus_bits)
        b_residues = self._reduce(b_bit_lists, modulus_bits)
        product = self._multiply(a_residues, b_residues, modulus_bits)

        for result_index, input_index in enumerate(valid_indices):
            outputs[input_index] = [int(bit) for bit in product[result_index].long().tolist()]
        return outputs

    def max_batch_size(self) -> int:
        return 256

    def _step(
        self,
        state_bits: torch.Tensor,
        multiplicand_bits: torch.Tensor,
        modulus_bits: torch.Tensor,
        control: torch.Tensor,
    ) -> torch.Tensor:
        if self.model is None:
            raise RuntimeError("model is not loaded")
        features = torch.stack(
            [state_bits, multiplicand_bits, modulus_bits],
            dim=-1,
        )
        logits = self.model(features, control)
        return (torch.sigmoid(logits) > 0.5).float()

    def _reduce(
        self,
        bit_lists: list[list[int]],
        modulus_bits: torch.Tensor,
    ) -> torch.Tensor:
        batch_size = len(bit_lists)
        operand_width = max(len(bits) for bits in bit_lists)
        padded = torch.zeros(
            (batch_size, operand_width),
            dtype=torch.long,
            device=self.device,
        )
        for row, bits in enumerate(bit_lists):
            if bits:
                padded[row, operand_width - len(bits) :] = torch.tensor(
                    bits,
                    dtype=torch.long,
                    device=self.device,
                )
        state_bits = torch.zeros(
            (batch_size, self._sequence_width),
            device=self.device,
        )
        one_bits = to_bits_limbs(
            [1] * batch_size,
            self.device,
            self._sequence_width,
        ).float()
        for position in range(operand_width):
            state_bits = self._step(
                state_bits,
                one_bits,
                modulus_bits,
                padded[:, position],
            )
        return state_bits

    def _multiply(
        self,
        a_residue: torch.Tensor,
        b_residue: torch.Tensor,
        modulus_bits: torch.Tensor,
    ) -> torch.Tensor:
        batch_size = a_residue.shape[0]
        state_bits = torch.zeros(
            (batch_size, self._sequence_width),
            device=self.device,
        )
        b_digits = b_residue.long()
        for position in range(self._sequence_width):
            state_bits = self._step(
                state_bits,
                a_residue,
                modulus_bits,
                b_digits[:, position],
            )
        return state_bits