"""Smaller, lower-compute inference wrapper for NeuralHorner v8.

The learned transition is unchanged.  Compared with the published wrapper:

* checkpoint tensors may be stored in bfloat16 and are restored to float32;
* logits are thresholded at zero (exactly equivalent to sigmoid(logit) > 0.5);
* only one operand is reduced before multiplication.  The other operand is
  streamed directly through the same Horner transition, eliminating a full
  modulus-width recurrent pass.
"""

from __future__ import annotations

from pathlib import Path

import torch
from torch import nn

from modchallenge.interface.base_model import ModularMultiplicationModel

_MASK32 = (1 << 32) - 1


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


def to_bits_limbs(ints, dev, width: int) -> torch.Tensor:
    nl = (width + 31) // 32
    cols = []
    for k in range(nl - 1, -1, -1):
        limb = torch.tensor(
            [(v >> (32 * k)) & _MASK32 for v in ints],
            dtype=torch.int64,
            device=dev,
        )
        cols.append(_to_bits_small(limb, 32))
    bits = torch.cat(cols, dim=1)
    return bits[:, nl * 32 - width:] if width < nl * 32 else bits


class Cell(nn.Module):
    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, d):
        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]:
    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
        self.L = 32
        self._Leff = 32

    def load(self, model_dir: str) -> None:
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")
        ckpt = torch.load(
            Path(model_dir) / "weights.pt",
            map_location="cpu",
            weights_only=True,
        )
        self.L = int(ckpt.get("L", 32))
        self.model = Cell(**ckpt.get("config", {}))
        # load_state_dict casts compact bf16 checkpoint tensors back to fp32.
        self.model.load_state_dict(ckpt["state_dict"])
        self.model.to(self.device)
        self.model.eval()
        self.model.gru.flatten_parameters()

    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)

    @torch.inference_mode()
    def predict_digits(self, a_enc, b_enc, p_enc):
        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]

    @torch.inference_mode()
    def predict_digits_batch(self, inputs):
        L = self.L
        max_op = 4 * L
        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)
            if p < 2 or p >= (1 << L) or len(a_bits) > max_op or len(b_bits) > max_op:
                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
        maxp = max(int(p).bit_length() for p in p_vals)
        self._Leff = min(self.L, max(32, ((maxp + 31) // 32) * 32))
        p_bits = to_bits_limbs(p_vals, dev, self._Leff).float()

        # (a*b) mod p = ((a mod p)*b) mod p.  Streaming the original b bits
        # through the learned Horner cell avoids first reducing b and then
        # scanning its L-bit residue a second time.
        ra = self._reduce(a_lists, p_bits, dev)
        prod = self._scan(b_lists, ra, p_bits, dev)
        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

    def _step(self, s_bits, feat, d):
        # The multiplicand and modulus channels stay constant for an entire
        # scan. Reuse their preallocated feature tensor instead of rebuilding
        # and copying all three channels at every recurrent step.
        feat[:, :, 0].copy_(s_bits)
        if self.device is not None and self.device.type == "cuda":
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                logits = self.model(feat, d)
            # Comparing a bf16 value with zero has the same sign decision as
            # first widening it to fp32, without allocating the fp32 logits.
            return (logits > 0).float()
        return (self.model(feat, d) > 0).float()

    def _scan(self, bit_lists, x_bits, p_bits, dev):
        n = len(bit_lists)
        width = max(len(bits) for bits in bit_lists)
        padded = torch.zeros((n, width), dtype=torch.long, device=dev)
        for row, bits in enumerate(bit_lists):
            if bits:
                padded[row, width - len(bits):] = torch.tensor(
                    bits, dtype=torch.long, device=dev
                )
        state = torch.zeros((n, self._Leff), device=dev)
        feat = torch.empty((n, self._Leff, 3), device=dev)
        feat[:, :, 1].copy_(x_bits)
        feat[:, :, 2].copy_(p_bits)
        for pos in range(width):
            state = self._step(state, feat, padded[:, pos])
        return state

    def _reduce(self, bit_lists, p_bits, dev):
        ones = to_bits_limbs([1] * len(bit_lists), dev, self._Leff).float()
        return self._scan(bit_lists, ones, p_bits, dev)