from __future__ import annotations

from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

from modchallenge.interface.base_model import ModularMultiplicationModel


MAX_P_BITS = 32          
PAD_HEAD = 3               

REDUCE_FEATURES = 6       
ADD_FEATURES = 5          


def _bits_of(n: int) -> list[int]:
    return [int(c) for c in bin(n)[2:]]

def linear_scan(alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:

    A, B = alpha, beta
    n = A.shape[1]
    off = 1
    while off < n:
        a_prev = F.pad(A, (0, 0, off, 0), value=1.0)[:, :n]
        b_prev = F.pad(B, (0, 0, off, 0), value=0.0)[:, :n]
        B = A * b_prev + B
        A = A * a_prev
        off <<= 1
    return B


class GateFn(nn.Module):

    def __init__(self, mode: str = "hard"):
        super().__init__()
        self.mode = mode

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        if self.mode == "soft":
            return torch.sigmoid(z)
        hard = (z > 0).to(z.dtype)
        if self.mode == "ste" and self.training:
            soft = torch.sigmoid(z)
            return hard + soft - soft.detach()
        return hard


class BiScanBlock(nn.Module):

    def __init__(self, d_model: int, d_scan: int, gate: GateFn):
        super().__init__()
        self.gate = gate
        self.proj_f = nn.Linear(d_model, 2 * d_scan)
        self.proj_b = nn.Linear(d_model, 2 * d_scan)
        self.out = nn.Linear(2 * d_scan, d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 2 * d_model),
            nn.ReLU(),
            nn.Linear(2 * d_model, d_model),
        )
        self.scan_noise = 0.0  

    def forward(self, u: torch.Tensor) -> torch.Tensor:
        zf = self.proj_f(u)
        zb = self.proj_b(u)
        af, bf = zf.chunk(2, dim=-1)
        ab, bb = zb.chunk(2, dim=-1)
        hf = linear_scan(self.gate(af), bf)
        hb = linear_scan(self.gate(ab.flip(1)), bb.flip(1)).flip(1)
        h = torch.cat([hf, hb], dim=-1)
        if self.training:
            self.last_h_l1 = h.abs().mean()
        if self.training and self.scan_noise > 0:
            h = h + torch.randn_like(h) * self.scan_noise
        u = u + self.out(h)
        u = u + self.mlp(u)
        return u


class BitCell(nn.Module):

    def __init__(self, n_features: int, n_borrow: int, n_q: int,
                 d_model: int = 32, d_scan: int = 16, n_blocks: int = 3,
                 gate_mode: str = "hard"):
        super().__init__()
        self.gate = GateFn(gate_mode)
        self.embed = nn.Linear(n_features, d_model)
        self.pre_mlp = nn.Sequential(
            nn.Linear(d_model, 2 * d_model),
            nn.ReLU(),
            nn.Linear(2 * d_model, d_model),
        )
        self.blocks = nn.ModuleList(
            BiScanBlock(d_model, d_scan, self.gate) for _ in range(n_blocks)
        )
        self.head = nn.Linear(d_model, 1)
        self.head_carry = nn.Linear(d_model, 1)
        self.head_sum = nn.Linear(d_model, 1)
        self.head_borrow = nn.Linear(d_model, n_borrow)
        self.head_q = nn.Linear(d_model, n_q)
        self.config = dict(n_features=n_features, n_borrow=n_borrow, n_q=n_q,
                           d_model=d_model, d_scan=d_scan, n_blocks=n_blocks)

    def trunk(self, feats: torch.Tensor):
        u = self.embed(feats)
        u = u + self.pre_mlp(u)
        taps = []
        for blk in self.blocks:
            u = blk(u)
            taps.append(u)
        return u, taps

    def forward(self, feats: torch.Tensor) -> torch.Tensor:
        u, _ = self.trunk(feats)
        return self.head(u).squeeze(-1)

    def forward_train(self, feats: torch.Tensor):
        u, taps = self.trunk(feats)
        return {
            "bits": self.head(u).squeeze(-1),
            "carry": self.head_carry(taps[0]).squeeze(-1),
            "sum": self.head_sum(taps[0]).squeeze(-1),
            "borrow": self.head_borrow(taps[1]),
            "q": self.head_q(taps[-1].mean(dim=1)),
        }


def make_reduce_cell(gate_mode: str = "hard", **kw) -> BitCell:
    return BitCell(REDUCE_FEATURES, n_borrow=3, n_q=4,
                   gate_mode=gate_mode, **kw)


def make_add_cell(gate_mode: str = "hard", **kw) -> BitCell:
    kw.setdefault("n_blocks", 2)  
    return BitCell(ADD_FEATURES, n_borrow=1, n_q=2, gate_mode=gate_mode, **kw)

def shift_bits(t: torch.Tensor, k: int) -> torch.Tensor:

    if k == 0:
        return t
    return torch.cat([t[:, k:], t.new_zeros(t.shape[0], k)], dim=1)


def _flags(B: int, N: int, dev) -> tuple[torch.Tensor, torch.Tensor]:
    is_msb = torch.zeros(B, N, device=dev)
    is_msb[:, 0] = 1.0
    is_lsb = torch.zeros(B, N, device=dev)
    is_lsb[:, -1] = 1.0
    return is_msb, is_lsb


def reduce_features(x, p, p3):
    B, N = x.shape
    is_msb, is_lsb = _flags(B, N, x.device)
    return torch.stack([x, p, shift_bits(p, 1), p3, is_msb, is_lsb], dim=-1)


def add_features(x, y, g):
    B, N = x.shape
    is_msb, is_lsb = _flags(B, N, x.device)
    return torch.stack([x, y, g.unsqueeze(1).expand(B, N),
                        is_msb, is_lsb], dim=-1)


class BitStreamMachine:

    def __init__(self, reduce_cell: BitCell, add_cell: BitCell,
                 device: torch.device):
        self.reduce_cell = reduce_cell
        self.add_cell = add_cell
        self.device = device

    @torch.no_grad()
    def _rstep(self, x, p, p3):
        logits = self.reduce_cell(reduce_features(x, p, p3))
        return (logits > 0).to(x.dtype)

    @torch.no_grad()
    def _astep(self, x, y, g):
        logits = self.add_cell(add_features(x, y, g))
        return (logits > 0).to(x.dtype)

    @torch.no_grad()
    def run(self, a_bits, b_bits, p_bits, p3_bits):

        B, N = p_bits.shape
        L = a_bits.shape[1]

        ops = torch.cat([a_bits, b_bits], dim=0)
        p2r = torch.cat([p_bits, p_bits], dim=0)
        p32r = torch.cat([p3_bits, p3_bits], dim=0)
        X = p2r.new_zeros(2 * B, N)
        for t in range(0, L, 2):
            x = torch.cat([X[:, 2:], ops[:, t: t + 2]], dim=1)
            X = self._rstep(x, p2r, p32r)
        ra, rb = X[:B], X[B:]

        Z = p_bits.new_zeros(B, N)
        for t in range(PAD_HEAD, N):
            s = self._astep(shift_bits(Z, 1), rb, ra[:, t])
            Z = self._rstep(s, p_bits, p3_bits)
        return Z

class BitStreamModel(ModularMultiplicationModel):
    def __init__(self):
        self.machine: BitStreamMachine | None = None
        self.device = torch.device("cpu") 

    def load(self, model_dir: str) -> None:
        torch.set_grad_enabled(False)
        ckpt = torch.load(
            Path(model_dir) / "weights.pt",
            map_location="cpu",
            weights_only=True,
        )
        rcell = make_reduce_cell()
        rcell.load_state_dict(ckpt["reduce_state_dict"], strict=True)
        rcell.eval()
        acell = make_add_cell()
        acell.load_state_dict(ckpt["add_state_dict"], strict=True)
        acell.eval()
        self.machine = BitStreamMachine(rcell, acell, self.device)


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

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

    def preprocess_p(self, p: str):
        v = int(p)
        return {"p": _bits_of(v), "p3": _bits_of((v << 1) + v)}


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

    @torch.no_grad()
    def predict_digits_batch(self, inputs) -> list[list[int]]:
        out: list[list[int] | None] = [None] * len(inputs)

        idx = [i for i, (_, _, pe) in enumerate(inputs)
               if len(pe["p"]) <= MAX_P_BITS]
        keep = set(idx)
        for i in range(len(inputs)):
            if i not in keep:
                out[i] = [0]
        if not idx:
            return [o if o is not None else [0] for o in out]

        sub = [inputs[i] for i in idx]
        n_p = max(len(pe["p"]) for _, _, pe in sub) + PAD_HEAD
        L = max(2, max(max(len(ae), len(be)) for ae, be, _ in sub))
        L += L % 2

        def pack(rows: list[list[int]], width: int) -> torch.Tensor:
            t = torch.zeros(len(rows), width)
            for r, bits in enumerate(rows):
                if bits:
                    t[r, width - len(bits):] = torch.tensor(
                        bits, dtype=torch.float32)
            return t

        a_t = pack([ae for ae, _, _ in sub], L)
        b_t = pack([be for _, be, _ in sub], L)
        p_t = pack([pe["p"] for _, _, pe in sub], n_p)
        p3_t = pack([pe["p3"] for _, _, pe in sub], n_p)

        z = self.machine.run(a_t, b_t, p_t, p3_t)
        z_int = z.to(torch.int64).tolist()
        for row, i in enumerate(idx):
            out[i] = [int(v) for v in z_int[row]]
        return [o if o is not None else [0] for o in out]

    def max_batch_size(self) -> int:
        return 128