"""A hyperdimensional full adder.

The integer ALU used to decode a register to a Python ``int``, add, and
re-encode.  This module removes that step for ``ADD`` and ``SUB``: the sum is
computed by a ripple-carry circuit whose primitive is *binding plus cleanup*,
which is the only thing a hypervector machine has.

The circuit
-----------
A number is a superposition of bit-position atoms (``int_to_hv``).  Reading a
bit is a similarity probe; that is unavoidable, and it is what a bit line is
in hardware:

    bit_i(n)  ==  sim(hv, bind(NUM, perm(BITBASE, i))) >= CODEC_THRESHOLD

A full-adder row is the classic 8-entry truth table, and each row is stored as
a record in HDC storage:

    row = bundle([
        bind(role("key"),   key_of(a, b, c)),
        bind(role("sum"),   atom("s", a ^ b ^ c)),
        bind(role("carry"), atom("k", majority(a, b, c))),
    ])

Evaluating one stage is a *content-addressable lookup*: form the query key,
and ask storage which row it resembles.  Each stage reports the similarity it
resolved at, so a marginal carry chain is visible rather than assumed.

Why the key needs one atom pair **per slot**
-------------------------------------------
XOR is associative and commutative, so a key built as
``bind(bind(atom_a, atom_b), atom_c)`` only ever sees
``atom_a ^ atom_b ^ atom_c``.  With a single shared pair for zero and one that
is *parity*: ``1^1^1`` and ``1^0^0`` both collapse to the same vector, and the
8-row table degenerates to 2 distinguishable entries.  Tagging each input with
a role before the XOR does not help, for the same reason — the role constants
cancel into a fixed mask.

Giving each slot its own atom pair (``gate/atom/a/0``, ``gate/atom/a/1``, ...)
makes the eight keys XOR combinations of *independently chosen* constants, so
they are pairwise near-orthogonal and cleanup resolves them exactly.  That is
the same distinction as using distinct address lines rather than one shared
bit line.

Subtraction is addition of the complement with carry-in 1 — two's complement,
so there is no sign special case anywhere: ``a - b == a + ~b + 1`` and the
32-bit wrap is whatever falls out of dropping the final carry.

Cost: 32 stages x (8 similarity probes + 4 atom probes) at d=10 000 is roughly
0.9 ms per operation, versus ~0.05 ms for the old decode-and-add.  That is the
price of keeping the arithmetic in the vector domain, and it is stated in the
README rather than hidden.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Sequence

from .hypervec import (
    CODEC_THRESHOLD,
    DEFAULT_DIM,
    MAX_INT_BITS,
    bind,
    bundle,
    hvc,
    hv_bits_of_int,
    hv_from_bits,
    hv_to_bits,
    hv_to_int,
    role,
    sim,
)
from .memory import HDCStorage

PREFIX = "adder/row/"

#: Gate slots: two inputs, the carry in, then the sum and carry out.
GATE_SLOTS = ("a", "b", "c", "s", "k")


class AdderError(Exception):
    """The adder could not resolve a truth-table row."""


def wrap(value: int, width: int = MAX_INT_BITS) -> int:
    """Wrap an integer into a ``width``-bit two's complement range."""
    span = 1 << width
    return (value + (span >> 1)) % span - (span >> 1)


@dataclass(frozen=True)
class Stage:
    """One bit slice of the carry chain, as executed."""

    index: int
    a: int
    b: int
    carry_in: int
    total: int
    carry_out: int
    confidence: float

    def render(self) -> str:
        return (
            "bit %2d  a=%d b=%d cin=%d -> sum=%d cout=%d  (lookup sim %.3f)"
            % (self.index, self.a, self.b, self.carry_in, self.total, self.carry_out,
               self.confidence)
        )


@dataclass(frozen=True)
class Row:
    """One truth-table row, and the record that stores it."""

    index: int
    a: int
    b: int
    c: int
    total: int
    carry: int
    key: int
    cell: str

    def render(self) -> str:
        return "a=%d b=%d cin=%d | sum=%d carry=%d | %s" % (
            self.a, self.b, self.c, self.total, self.carry, self.cell
        )


class HyperAdder:
    """Ripple-carry adder whose gates are HDC bindings and cleanups."""

    def __init__(
        self,
        d: int = DEFAULT_DIM,
        storage: HDCStorage | None = None,
        width: int = MAX_INT_BITS,
        threshold: float = CODEC_THRESHOLD,
        radix: int = 1,
    ):
        self.d = d
        self.width = min(width, MAX_INT_BITS)
        self.threshold = threshold
        #: Bits per stage.  1 is the 8-row full adder (32 stages); 8 is a
        #: 131 072-row byte gate (4 stages), materialised lazily.
        self.radix = max(1, min(radix, 8))
        #: One atom pair per slot — see the module docstring on why sharing
        #: a single pair would collapse the key to parity.
        self.atoms = {
            (slot, bit): hvc("gate/atom/%s/%d" % (slot, bit), d)
            for slot in GATE_SLOTS
            for bit in (0, 1)
        }
        self.storage = storage if storage is not None else HDCStorage(
            d=d, capacity=64, label="adder"
        )
        self.rows: list[Row] = []
        self.last_stages: list = []
        self.stages_executed = 0
        self.gate = None
        if self.radix == 1:
            self._build_table()
        else:
            from .slicegate import SliceGate

            self.gate = SliceGate(radix=self.radix, d=d, threshold=threshold)

    # ------------------------------------------------------------ the gate
    def atom(self, slot: str, bit: int) -> int:
        """The vector that represents ``bit`` in a given gate slot."""
        return self.atoms[(slot, 1 if bit else 0)]

    def key_of(self, a: int, b: int, c: int) -> int:
        """Exact gate address: XOR of three role-specific atoms."""
        return bind(
            bind(self.atom("a", a), self.atom("b", b), self.d),
            self.atom("c", c),
            self.d,
        )

    def _build_table(self) -> None:
        """Materialise the 8 full-adder rows as cells in HDC storage."""
        for a in (0, 1):
            for b in (0, 1):
                for c in (0, 1):
                    total = a ^ b ^ c
                    carry = 1 if (a + b + c) >= 2 else 0
                    key = self.key_of(a, b, c)
                    record = bundle(
                        [
                            bind(role("key", self.d), key, self.d),
                            bind(role("sum", self.d), self.atom("s", total), self.d),
                            bind(role("carry", self.d), self.atom("k", carry), self.d),
                        ],
                        self.d,
                    )
                    cell = PREFIX + "%d%d%d" % (a, b, c)
                    self.storage.write(cell, record, symbol=True)
                    self.rows.append(
                        Row(a * 4 + b * 2 + c, a, b, c, total, carry, key, cell)
                    )

    def _pick(self, probe: int, slot: str) -> tuple[int, float]:
        """Cleanup an output field to 0 or 1; returns (bit, margin)."""
        s0 = sim(probe, self.atom(slot, 0), self.d)
        s1 = sim(probe, self.atom(slot, 1), self.d)
        if s1 >= s0:
            return 1, s1 - s0
        return 0, s0 - s1

    def full_adder(self, a: int, b: int, c: int) -> tuple[int, int, float]:
        """One stage: bind the inputs, clean up over the 8 stored rows."""
        query = self.key_of(a, b, c)
        # The stored row holds bind(role("key"), key), so the probe carries the
        # same role; a bare key against the bundle would sit at the noise floor.
        probe = bind(role("key", self.d), query, self.d)
        hits = self.storage.recall(probe, k=1, prefix=PREFIX)
        if not hits or hits[0][2] < 0.20:
            raise AdderError(
                "no truth-table row resolved (best similarity %.3f)"
                % (hits[0][2] if hits else -1.0)
            )
        record, confidence = hits[0][1], hits[0][2]
        total, _ = self._pick(bind(record, role("sum", self.d), self.d), "s")
        carry, _ = self._pick(bind(record, role("carry", self.d), self.d), "k")
        return total, carry, confidence

    # --------------------------------------------------------- the ripple
    def add_bits(
        self, a_bits: Sequence[int], b_bits: Sequence[int], carry_in: int = 0
    ) -> tuple[list[int], list]:
        """Carry across the operand, one gate stage at a time."""
        if self.gate is not None:
            return self._add_sliced(a_bits, b_bits, carry_in)
        carry = 1 if carry_in else 0
        out: list[int] = []
        stages: list[Stage] = []
        for i in range(self.width):
            # A short slice is zero-extended, exactly as a partial bit vector
            # is on a real adder input.
            a = 1 if i < len(a_bits) and a_bits[i] else 0
            b = 1 if i < len(b_bits) and b_bits[i] else 0
            total, carry_next, confidence = self.full_adder(a, b, carry)
            stages.append(Stage(i, a, b, carry, total, carry_next, confidence))
            out.append(total)
            carry = carry_next
        self.stages_executed += len(stages)
        self.last_stages = stages
        return out, stages

    def _add_sliced(
        self, a_bits: Sequence[int], b_bits: Sequence[int], carry_in: int
    ) -> tuple[list[int], list]:
        """Width-k stages: each gate call consumes ``radix`` bits plus carry."""
        from .slicegate import SliceStep

        k = self.radix
        out = [0] * self.width
        steps: list[SliceStep] = []
        carry = 1 if carry_in else 0
        for start in range(0, self.width, k):
            a_val = 0
            b_val = 0
            for i in range(k):
                pos = start + i
                if pos < self.width:
                    a_val |= (1 if pos < len(a_bits) and a_bits[pos] else 0) << i
                    b_val |= (1 if pos < len(b_bits) and b_bits[pos] else 0) << i
            total, carry_next, confidence = self.gate.evaluate(a_val, b_val, carry)
            for i in range(k):
                pos = start + i
                if pos < self.width:
                    out[pos] = (total >> i) & 1
            steps.append(
                SliceStep(start, a_val, b_val, carry, total, carry_next, confidence, k)
            )
            carry = carry_next
        self.stages_executed += len(steps)
        self.last_stages = steps
        return out, steps

    # -------------------------------------------------------- public API
    def add(self, a: int, b: int) -> int:
        """Integer operands in, hypervector out."""
        bits, _ = self.add_bits(hv_bits_of_int(a), hv_bits_of_int(b), 0)
        return hv_from_bits(bits, self.d)

    def sub(self, a: int, b: int) -> int:
        """``a - b`` as ``a + ~b + 1``, with no sign case analysis."""
        inverted = [0 if bit else 1 for bit in hv_bits_of_int(b)]
        bits, _ = self.add_bits(hv_bits_of_int(a), inverted, 1)
        return hv_from_bits(bits, self.d)

    def add_hv(self, a_hv: int, b_hv: int) -> int:
        """Hypervector in, hypervector out — no integer is materialised."""
        bits, _ = self.add_bits(
            hv_to_bits(a_hv, self.d, self.threshold),
            hv_to_bits(b_hv, self.d, self.threshold),
            0,
        )
        return hv_from_bits(bits, self.d)

    def sub_hv(self, a_hv: int, b_hv: int) -> int:
        """As :meth:`add_hv`, via the complement and carry-in 1."""
        b_bits = hv_to_bits(b_hv, self.d, self.threshold)
        bits, _ = self.add_bits(
            hv_to_bits(a_hv, self.d, self.threshold),
            [0 if bit else 1 for bit in b_bits],
            1,
        )
        return hv_from_bits(bits, self.d)

    def decode(self, hv: int) -> int:
        """Decode to an integer — for display and for verifying the circuit."""
        return hv_to_int(hv, self.d, self.threshold)

    # ------------------------------------------------------------- views
    def table_text(self) -> str:
        rows = ["full-adder truth table, stored as records in %r" % self.storage.label]
        for row in self.rows:
            rows.append("  " + row.render())
        return "\n".join(rows)

    def trace_text(self, limit: int | None = None) -> str:
        stages = self.last_stages if limit is None else self.last_stages[:limit]
        if not stages:
            return "<no stages executed yet>"
        return "\n".join(stage.render() for stage in stages)

    def separation(self) -> dict:
        """Measured similarity between the 8 row keys (signal vs crosstalk)."""
        if self.gate is not None:
            summary = self.gate.summary()
            summary.update({
                "rows": summary["table_rows"],
                "worst_cross_similarity": None,
                "table_cells": summary["cells"],
                "width": self.width,
                "d": self.d,
            })
            return summary
        worst = 2.0
        for i, row_a in enumerate(self.rows):
            for j, row_b in enumerate(self.rows):
                if i != j:
                    worst = min(worst, sim(row_a.key, row_b.key, self.d))
        return {
            "rows": len(self.rows),
            "worst_cross_similarity": worst,
            "table_cells": len(self.storage),
            "width": self.width,
            "d": self.d,
        }

    def verify(self, trials: int = 200, seed: int = 1234) -> dict:
        """Check the circuit against Python arithmetic, both directions."""
        import random

        rng = random.Random(seed)
        adds = subs = hv_layer = 0
        for _ in range(trials):
            a = rng.randrange(-(1 << 31), 1 << 31)
            b = rng.randrange(-(1 << 31), 1 << 31)
            if self.decode(self.add(a, b)) == wrap(a + b, self.width):
                adds += 1
            if self.decode(self.sub(a, b)) == wrap(a - b, self.width):
                subs += 1
            if self.decode(self.add_hv(self.add(0, a), self.add(0, b))) == wrap(
                a + b, self.width
            ):
                hv_layer += 1
        return {
            "trials": trials,
            "add_correct": adds,
            "sub_correct": subs,
            "hv_to_hv_correct": hv_layer,
            "stages": self.width,
        }


def selfcheck(trials: int = 100) -> dict:
    """Small self-verification used by the demo and the CLI."""
    adder = HyperAdder(d=2000)
    report = adder.separation()
    report.update(adder.verify(trials))
    report["decode_roundtrip"] = all(
        adder.decode(adder.add(0, n)) == n for n in (0, 1, -1, 7, -7, 65535, -12345)
    )
    return report


if __name__ == "__main__":  # pragma: no cover - manual inspection
    for key, value in selfcheck().items():
        print("%-22s %s" % (key, value))