"""Constructed ReLU circuit for exact (a * b) mod p.

This module separates two concerns on purpose:

* **Topology** — which ReLU/linear/conv primitives run, in what order, and at
  what tier geometry (limb count ``n``, derived from the prime's bit width).
  Built by :func:`build_topology`; it is the same object the frontier-track
  training experiment will later optimise from random init.
* **Weight filling** — the numeric content of every primitive (step
  thresholds, gating constants, shift amounts). Supplied by an *initializer*.
  :class:`ConstructedInit` fills the bit-exact arithmetic-circuit weights;
  :class:`RandomInit` fills the same slots from a seeded RNG so the identical
  topology is a trainable network. The constructed weights are therefore one
  initializer among several, not a baked-in forward pass.

Forward-pass discipline (the competition's letter): every operation is a
linear map, a ReLU, or a 1-D convolution. There is no integer tensor
arithmetic, no ``einsum`` on the inputs, and no product of two activations.
Multiplication is done by the binary-gated-product identity, which replaces
every bilinear op with a ReLU. Weights are all in ``{0, +-1, +-2^t}`` and are
exactly representable in float32 storage; compute runs in float64.

The three identities (each exact on integer-valued floats):

1. step gate: ``step_tau(x) = relu(x - tau + 1) - relu(x - tau)`` is ``1{x >= tau}``.
2. gated product: for ``0 <= v < 2^16`` and bit ``b``, ``b*v = relu(v - 2^16*(1-b))``.
3. boolean: ``AND = relu(a+b-1)``, ``XOR = a+b-2*AND``, ``OR = a+b-AND``.

See ``src/constructed/README.md`` for the topology/initializer split and the
companion feasibility spec for the per-tier envelope.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.nn as nn

LIMB_BITS = 16
B = 1 << LIMB_BITS  # base 2^16


def n_limbs_for_bits(max_bits: int) -> int:
    """Limb count for a prime up to ``max_bits`` wide, base 2^16."""
    return max(1, math.ceil(max_bits / LIMB_BITS))


# ---------------------------------------------------------------------------
# ReLU-only primitives. Each is a free function on float64 tensors so the same
# code path serves the constructed weights and (later) trained ones; the only
# learnable content is the constants passed in, surfaced by the initializer.
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class Consts:
    """The learnable constants threaded through the primitives.

    ``one`` is the step-gate unit width (constructed: ``1.0``); ``base`` is the
    gated-product gating constant (constructed: ``2^16``). Passing them as data
    rather than hardcoding makes the topology genuinely trainable and makes the
    weight-perturbation collapse test bite: a wrong ``one`` or ``base`` breaks
    every comparator and gated product in the circuit.
    """

    one: float = 1.0
    base: float = float(B)


_CONSTRUCTED = Consts()


class _GateCounter:
    """Opt-in tally of ReLU gates and circuit depth, for the per-tier audit.

    Counting is a measurement aid only; it never changes the forward result.
    A "gate" is one ReLU lane (per scalar position), accumulated across the
    sequential pipeline so the number is comparable to the doc's ReLUs/prob.
    """

    def __init__(self):
        self.active = False
        self.gates = 0
        self.relu_calls = 0

    def add(self, t: torch.Tensor) -> None:
        if not self.active:
            return
        self.relu_calls += 1
        # lanes per problem = elements / batch (last-dim positions per item)
        self.gates += t.shape[-1] if t.dim() else 1


_COUNTER = _GateCounter()


def relu(t: torch.Tensor) -> torch.Tensor:
    _COUNTER.add(t)
    return torch.clamp(t, min=0.0)


class count_gates:
    """Context manager: tally ReLU gates/calls during the enclosed forward.

    >>> with count_gates() as c:
    ...     circuit(x_limbs, y_bits, p_bits, mu_bits)
    >>> c.gates, c.relu_calls
    """

    def __enter__(self) -> _GateCounter:
        _COUNTER.active = True
        _COUNTER.gates = 0
        _COUNTER.relu_calls = 0
        return _COUNTER

    def __exit__(self, *exc) -> None:
        _COUNTER.active = False


def step_gate(x: torch.Tensor, tau: float, one: float = 1.0) -> torch.Tensor:
    """Exact ``1{x >= tau}`` on integer-valued ``x`` via two ReLUs (identity 1).

    ``one`` is the unit-step width: the constructed value is ``1.0``. It is
    surfaced as a learnable constant so the topology is trainable and so a
    perturbed value provably breaks the comparator (the operational test).
    """
    return relu(x - tau + one) - relu(x - tau)


def gated_product(v: torch.Tensor, b: torch.Tensor, base: float = float(B)) -> torch.Tensor:
    """Exact ``b*v`` for bit ``b`` and ``0 <= v < 2^16`` via one ReLU (identity 2).

    ``base`` is the gating constant ``2^16``; surfaced as a learnable constant
    for the same reason as ``one`` above.
    """
    return relu(v - base * (1.0 - b))


def bit_and(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return relu(a + b - 1.0)


def bit_xor(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a + b - 2.0 * bit_and(a, b)


def peel_bits(
    value: torch.Tensor, width: int, consts: Consts = _CONSTRUCTED
) -> torch.Tensor:
    """MSB-first sequential bit peel of an integer in ``[0, 2^width)``.

    Returns a ``(..., width)`` tensor of bits, LSB first. Each bit is a
    2-ReLU step comparator at a power-of-2 threshold; the residual update
    ``r <- r - b*2^i`` is linear because ``b`` is already exactly ``0/1``.
    """
    r = value.clone()
    bits = []
    for i in range(width - 1, -1, -1):
        b = step_gate(r, float(1 << i), consts.one)
        r = r - b * float(1 << i)
        bits.append(b)
    bits.reverse()
    return torch.stack(bits, dim=-1)


def limbs_from_columns(
    cols: torch.Tensor, col_width: int, consts: Consts = _CONSTRUCTED
) -> torch.Tensor:
    """Carry-propagate a base-2^16 carry-save column vector into clean limbs.

    ``cols`` is ``(..., m)`` of nonnegative integer-valued sums, each
    ``< 2^col_width``. Returns ``(..., m + extra)`` base-2^16 limbs of the same
    total value. The ripple uses :func:`peel_bits` per column (the doc's bit
    peel + re-bin), splitting each running sum into a low 16-bit limb and the
    carry into the next column. Pure linear + ReLU.
    """
    m = cols.shape[-1]
    lead = cols.shape[:-1]
    carry = torch.zeros(lead, dtype=cols.dtype)
    out = []

    def split(s: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        bits = peel_bits(s, col_width, consts)
        weights_lo = torch.tensor(
            [float(1 << j) for j in range(LIMB_BITS)], dtype=cols.dtype
        )
        weights_hi = torch.tensor(
            [float(1 << (j - LIMB_BITS)) for j in range(LIMB_BITS, col_width)],
            dtype=cols.dtype,
        )
        lo = (bits[..., :LIMB_BITS] * weights_lo).sum(dim=-1)
        hi = (bits[..., LIMB_BITS:] * weights_hi).sum(dim=-1)
        return lo, hi

    for i in range(m):
        s = cols[..., i] + carry
        lo, carry = split(s)
        out.append(lo)
    # flush remaining carry (bounded; carry shrinks each step)
    while bool((carry > 0).any()):
        lo, carry = split(carry)
        out.append(lo)
    return torch.stack(out, dim=-1)


# ---------------------------------------------------------------------------
# Bignum operations spelled with the primitives above. A bignum is a
# (..., L) tensor of base-2^16 limb values, little-endian. The MULTIPLIER is
# always supplied as bits (the gating operand), never as a value, so the
# bilinear "value times value" never appears.
# ---------------------------------------------------------------------------

def schoolbook_mul(
    x_limbs: torch.Tensor,
    y_bits: torch.Tensor,
    peak: list | None = None,
    consts: Consts = _CONSTRUCTED,
) -> torch.Tensor:
    """``x * y`` where ``x`` is base-2^16 limbs and ``y`` is bits (LSB first).

    Gated partial products (identity 2) accumulated into carry-save columns,
    then carry-propagated. Each column sum stays ``< n * 2^32`` (the doc's
    bound), inside float64's exact range with a wide margin. If ``peak`` is a
    list, the maximum carry-save column magnitude seen is appended to it (the
    largest intermediate of the whole circuit, for the precision audit).
    """
    n = x_limbs.shape[-1]
    nbits = y_bits.shape[-1]
    lead = x_limbs.shape[:-1]
    max_col = n + (nbits + LIMB_BITS - 1) // LIMB_BITS + 2
    col_width = (2 * LIMB_BITS) + max(1, (n).bit_length() + nbits.bit_length())
    cols = [torch.zeros(lead, dtype=x_limbs.dtype) for _ in range(max_col)]
    for j in range(nbits):
        bj = y_bits[..., j]
        off = j % LIMB_BITS
        base_col = j // LIMB_BITS
        scale = float(1 << off)
        for i in range(n):
            gx = gated_product(x_limbs[..., i], bj, consts.base)  # bj * x_limbs[i]
            cols[base_col + i] = cols[base_col + i] + gx * scale
    stacked = torch.stack(cols, dim=-1)
    if peak is not None and stacked.numel():
        peak.append(float(stacked.abs().max()))
    return limbs_from_columns(stacked, col_width, consts)


def int_to_limbs(value: int, n: int) -> list[int]:
    return [(value >> (LIMB_BITS * i)) & (B - 1) for i in range(n)]


def int_to_bits(value: int, nbits: int) -> list[int]:
    return [(value >> i) & 1 for i in range(nbits)]


def limbs_to_int(limbs: torch.Tensor) -> int:
    vals = [int(round(float(v))) for v in limbs.flatten().tolist()]
    return sum(v << (LIMB_BITS * i) for i, v in enumerate(vals))


def bits_to_int(bits: torch.Tensor) -> int:
    vals = [int(round(float(v))) for v in bits.flatten().tolist()]
    return sum(v << i for i, v in enumerate(vals))


# ---------------------------------------------------------------------------
# Topology + initializers
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class TierGeometry:
    """Geometry the circuit is built for; derived from the prime bit width."""

    tier_idx: int
    max_bits: int

    @property
    def n(self) -> int:
        return n_limbs_for_bits(self.max_bits)

    @property
    def k(self) -> int:
        # Barrett uses k = n limbs (p < 2^16k = b^k).
        return self.n


@dataclass(frozen=True)
class Topology:
    """The untrainable wiring: tier geometry plus the fixed structural plan.

    The frontier-track training experiment treats this as fixed and learns the
    initializer-filled constants. Nothing here depends on a particular (a,b,p).
    """

    geom: TierGeometry

    @property
    def n(self) -> int:
        return self.geom.n

    @property
    def k(self) -> int:
        return self.geom.k

    # column width for the n-limb * 2n-bit product accumulation
    @property
    def prod_col_width(self) -> int:
        n = self.n
        return (2 * LIMB_BITS) + max(1, (n).bit_length() + (2 * LIMB_BITS * n).bit_length())


def build_topology(tier_idx: int, max_bits: int) -> Topology:
    return Topology(TierGeometry(tier_idx=tier_idx, max_bits=max_bits))


class Initializer:
    """Fills the circuit's learnable constants. Constructed or random."""

    def step_one(self) -> float:
        raise NotImplementedError

    def gate_base(self) -> float:
        raise NotImplementedError


class ConstructedInit(Initializer):
    """Bit-exact arithmetic-circuit constants (the +-1 / 2^16 grid)."""

    def step_one(self) -> float:
        return 1.0

    def gate_base(self) -> float:
        return float(B)


class RandomInit(Initializer):
    """Same slots, seeded random content: the constructed weights become one
    initializer among several for the trainable topology."""

    def __init__(self, seed: int = 0):
        g = torch.Generator().manual_seed(seed)
        self._a = float(torch.rand((), generator=g))
        self._b = float(torch.rand((), generator=g)) * B

    def step_one(self) -> float:
        return self._a

    def gate_base(self) -> float:
        return self._b


# ---------------------------------------------------------------------------
# The circuit module
# ---------------------------------------------------------------------------

class ModmulCircuit(nn.Module):
    """Exact ``(a*b) mod p`` as a linear + ReLU circuit at a tier geometry.

    The constants the initializer fills are registered as float32 buffers so
    they save to safetensors on the exact weight grid; the forward pass casts
    to float64 for exact integer arithmetic. With :class:`ConstructedInit`
    these buffers are ``1.0`` and ``2^16``; randomising them collapses
    correctness, which is the operational test that the capability lives in the
    weights rather than in the wiring.

    Inputs to :meth:`forward` are the per-problem preprocessed tensors (limbs
    and bits), exactly what a competition ``predict_digits`` would compute
    outside the network. No (a,b,p) integer arithmetic happens in forward.
    """

    def __init__(self, topology: Topology, init: Initializer | None = None):
        super().__init__()
        self.topology = topology
        init = init or ConstructedInit()
        self.register_buffer(
            "step_one", torch.tensor(init.step_one(), dtype=torch.float32)
        )
        self.register_buffer(
            "gate_base", torch.tensor(init.gate_base(), dtype=torch.float32)
        )

    @property
    def n(self) -> int:
        return self.topology.n

    @property
    def k(self) -> int:
        return self.topology.k

    def _consts(self) -> Consts:
        """The learnable constants as a plain dataclass for the primitives."""
        return Consts(one=float(self.step_one), base=float(self.gate_base))

    # -- the staged pipeline -------------------------------------------------

    def _mul(
        self,
        x_limbs: torch.Tensor,
        y_bits: torch.Tensor,
        consts: Consts,
        peak: list | None = None,
    ) -> torch.Tensor:
        return schoolbook_mul(x_limbs, y_bits, peak=peak, consts=consts)

    def forward(
        self,
        x_limbs: torch.Tensor,
        y_bits: torch.Tensor,
        p_bits: torch.Tensor,
        mu_bits: torch.Tensor,
        trace: dict | None = None,
    ) -> torch.Tensor:
        """Compute base-2^16 limbs of ``(x*y) mod p``.

        Args:
            x_limbs: ``(..., n)`` base-2^16 limbs of ``x = a mod p`` in ``[0, p)``.
            y_bits:  ``(..., 16n)`` bits of ``y = b mod p`` (LSB first).
            p_bits:  ``(..., 16n)`` bits of ``p`` (LSB first).
            mu_bits: ``(..., 16(n+1)+1)`` bits of ``mu = floor(2^(32n)/p)``.

        Returns ``(..., n)`` limbs of the residue, all ``< 2^16``, value ``< p``.
        """
        x_limbs = x_limbs.double()
        y_bits = y_bits.double()
        p_bits = p_bits.double()
        mu_bits = mu_bits.double()
        n, k = self.n, self.k
        c = self._consts()
        peak_a: list = []
        peak_b: list = []
        peak_c: list = []

        # Stage A/N: t = x * y, 2n limbs.
        t = self._mul(x_limbs, y_bits, c, peak=peak_a)  # (..., <=2n+1)
        t = self._fit(t, 2 * n + 1)

        # Stage B: q1 = t >> 16(k-1)  (limb selection, free), then q-hat.
        q1 = t[..., (k - 1):]  # limbs of floor(t / b^(k-1))
        q1_bits = self._limbs_to_bits(q1, c)
        mu_lo = mu_bits[..., : (LIMB_BITS * (k + 1) + 1)]
        # q2 = q1 * mu ; gating operand = q1 bits (smaller side keeps cols tight)
        mu_limbs = self._bits_to_limbs(mu_lo)
        q2 = schoolbook_mul(mu_limbs, q1_bits, peak=peak_b, consts=c)
        # q3 = q2 >> 16(k+1)  (limb selection)
        q3 = q2[..., (k + 1):]
        q3 = self._fit(q3, n + 1)

        # Stage C: q3 * p, keep low (k+1) limbs (mod b^(k+1)).
        p_limbs = self._bits_to_limbs(p_bits)
        q3_bits = self._limbs_to_bits(q3, c)
        qp = schoolbook_mul(p_limbs, q3_bits, peak=peak_c, consts=c)
        qp_lo = self._fit(qp[..., : (k + 1)], k + 1)
        t_lo = self._fit(t[..., : (k + 1)], k + 1)

        # r = t_lo - qp_lo  (mod b^(k+1)); both nonneg, t_lo >= qp_lo here.
        r = self._sub_limbs(t_lo, qp_lo, k + 1, c)

        # Stage E: at most two conditional subtractions of p.
        p_lo = self._fit(p_limbs, k + 1)
        for _ in range(2):
            ge = self._ge(r, p_lo, k + 1, c)  # bit: r >= p ?
            r = self._cond_sub(r, p_lo, ge, k + 1, c)

        if trace is not None:
            trace["peak_mul_xy"] = max(peak_a, default=0.0)
            trace["peak_mul_q1mu"] = max(peak_b, default=0.0)
            trace["peak_mul_q3p"] = max(peak_c, default=0.0)
            trace["peak_overall"] = max(
                trace["peak_mul_xy"], trace["peak_mul_q1mu"], trace["peak_mul_q3p"]
            )
        return self._fit(r[..., :n], n)

    # -- helpers (all linear + ReLU) ----------------------------------------

    @staticmethod
    def _fit(limbs: torch.Tensor, width: int) -> torch.Tensor:
        """Pad/truncate a limb vector to ``width`` limbs (linear selection)."""
        cur = limbs.shape[-1]
        if cur == width:
            return limbs
        if cur > width:
            return limbs[..., :width]
        pad = torch.zeros(
            (*limbs.shape[:-1], width - cur), dtype=limbs.dtype
        )
        return torch.cat([limbs, pad], dim=-1)

    @staticmethod
    def _limbs_to_bits(limbs: torch.Tensor, consts: Consts) -> torch.Tensor:
        """Each base-2^16 limb -> 16 bits (LSB first), concatenated."""
        parts = [
            peel_bits(limbs[..., i], LIMB_BITS, consts)
            for i in range(limbs.shape[-1])
        ]
        return torch.cat(parts, dim=-1)

    @staticmethod
    def _bits_to_limbs(bits: torch.Tensor) -> torch.Tensor:
        """Group bits (LSB first) into base-2^16 limb values (linear)."""
        nbits = bits.shape[-1]
        nl = (nbits + LIMB_BITS - 1) // LIMB_BITS
        w = torch.tensor([float(1 << j) for j in range(LIMB_BITS)], dtype=bits.dtype)
        out = []
        for i in range(nl):
            chunk = bits[..., i * LIMB_BITS : (i + 1) * LIMB_BITS]
            if chunk.shape[-1] < LIMB_BITS:
                pad = torch.zeros(
                    (*chunk.shape[:-1], LIMB_BITS - chunk.shape[-1]), dtype=bits.dtype
                )
                chunk = torch.cat([chunk, pad], dim=-1)
            out.append((chunk * w).sum(dim=-1))
        return torch.stack(out, dim=-1)

    def _sub_limbs(
        self, a: torch.Tensor, b: torch.Tensor, width: int, consts: Consts
    ) -> torch.Tensor:
        """``a - b`` as base-2^16 limbs, assuming ``a >= b`` (borrow ripple).

        Borrow is a step comparator; the difference is recombined linearly.
        """
        a = self._fit(a, width)
        b = self._fit(b, width)
        borrow = torch.zeros(a.shape[:-1], dtype=a.dtype)
        out = []
        for i in range(width):
            d = a[..., i] - b[..., i] - borrow
            need = step_gate(-d, 1.0, consts.one)  # 1 if d < 0
            d = d + need * consts.base
            borrow = need
            out.append(d)
        return torch.stack(out, dim=-1)

    def _ge(
        self, a: torch.Tensor, b: torch.Tensor, width: int, consts: Consts
    ) -> torch.Tensor:
        """Bit ``1{a >= b}`` for base-2^16 limb vectors via borrow-out."""
        a = self._fit(a, width)
        b = self._fit(b, width)
        borrow = torch.zeros(a.shape[:-1], dtype=a.dtype)
        for i in range(width):
            d = a[..., i] - b[..., i] - borrow
            borrow = step_gate(-d, 1.0, consts.one)
        return 1.0 - borrow  # no final borrow => a >= b

    def _cond_sub(
        self,
        r: torch.Tensor,
        p: torch.Tensor,
        cond: torch.Tensor,
        width: int,
        consts: Consts,
    ) -> torch.Tensor:
        """``r - cond*p`` (cond a bit): gate p's limbs, then borrow-subtract."""
        p = self._fit(p, width)
        gated = torch.stack(
            [gated_product(p[..., i], cond, consts.base) for i in range(width)],
            dim=-1,
        )
        diff = self._sub_limbs(r, gated, width, consts)
        # when cond==0 the subtraction is a no-op; mux to keep r exactly
        keep = (1.0 - cond).unsqueeze(-1)
        take = cond.unsqueeze(-1)
        return self._fit(r, width) * keep + diff * take


# ---------------------------------------------------------------------------
# Preprocessing: turn integer (a, b, p) into the circuit's input tensors.
# This is caller-side representation work (base conversion + the p-derived
# constant mu), explicitly outside the forward pass.
# ---------------------------------------------------------------------------

def preprocess(a: int, b: int, p: int, geom: TierGeometry):
    """Build (x_limbs, y_bits, p_bits, mu_bits) for one problem.

    ``x = a mod p`` as limbs, ``y = b mod p`` as bits, ``p`` as bits, and
    ``mu = floor(2^(32n)/p)`` as bits. All integer work here is caller-side
    representation, never inside :meth:`ModmulCircuit.forward`.
    """
    n = geom.n
    x = a % p
    y = b % p
    mu = (1 << (2 * LIMB_BITS * n)) // p
    x_limbs = torch.tensor(int_to_limbs(x, n), dtype=torch.float64)
    y_bits = torch.tensor(int_to_bits(y, LIMB_BITS * n), dtype=torch.float64)
    p_bits = torch.tensor(int_to_bits(p, LIMB_BITS * n), dtype=torch.float64)
    mu_bits = torch.tensor(
        int_to_bits(mu, LIMB_BITS * (n + 1) + 1), dtype=torch.float64
    )
    return x_limbs, y_bits, p_bits, mu_bits


def preprocess_batch(triples, geom: TierGeometry):
    """Stack :func:`preprocess` over a list of ``(a, b, p)`` into batched tensors."""
    xl, yb, pb, mb = [], [], [], []
    for a, b, p in triples:
        a4, b4, c4, d4 = preprocess(a, b, p, geom)
        xl.append(a4)
        yb.append(b4)
        pb.append(c4)
        mb.append(d4)
    return (
        torch.stack(xl),
        torch.stack(yb),
        torch.stack(pb),
        torch.stack(mb),
    )


def run_one(circuit: ModmulCircuit, a: int, b: int, p: int) -> int:
    """Convenience: preprocess, forward, decode to an integer residue."""
    xl, yb, pb, mb = preprocess(a, b, p, circuit.topology.geom)
    out = circuit(xl, yb, pb, mb)
    return limbs_to_int(out)


# ---------------------------------------------------------------------------
# safetensors I/O — the constructed (or trained) weights as a portable artifact
# ---------------------------------------------------------------------------

def save_circuit(circuit: ModmulCircuit, path) -> None:
    """Write the circuit's filled constants to safetensors (float32 storage).

    The topology (tier id, max_bits) travels in the safetensors metadata so a
    load reconstructs the same wiring before refilling the weights.
    """
    from pathlib import Path

    from safetensors.torch import save_file

    path = Path(path)
    tensors = {k: v.contiguous() for k, v in circuit.state_dict().items()}
    meta = {
        "tier_idx": str(circuit.topology.geom.tier_idx),
        "max_bits": str(circuit.topology.geom.max_bits),
    }
    save_file(tensors, str(path), metadata=meta)


def load_circuit(path) -> ModmulCircuit:
    """Reconstruct topology from metadata, then load the stored constants."""
    from pathlib import Path

    from safetensors import safe_open
    from safetensors.torch import load_file

    path = Path(path)
    with safe_open(str(path), framework="pt") as f:
        meta = f.metadata() or {}
    topo = build_topology(int(meta["tier_idx"]), int(meta["max_bits"]))
    circuit = ModmulCircuit(topo)
    circuit.load_state_dict(load_file(str(path)))
    circuit.eval()
    return circuit