"""Wide slice gates: several bits of adder per stage.

The 1-bit gate in :mod:`hdc.adder` is 8 rows and 32 stages.  A gate that
handles ``k`` bits at once is ``2**(2k+1)`` rows and ``ceil(32/k)`` stages, so
the trade is explicit:

    k=1    8 rows        32 stages
    k=2    32 rows       16 stages
    k=4    512 rows       8 stages
    k=8    131072 rows    4 stages

Everything here is content-addressable in the same sense as the 1-bit gate: a
stage's inputs select a row *by vector resemblance*, and the row's outputs are
recovered by unbinding fields and cleaning them up against the gate atoms.

Two design points that are easy to get wrong
--------------------------------------------
**A lazy table.**  131 072 rows at d=10 000 would be 164 MB of vectors, and
enumerating them costs ~1.6 s of XOR folding before the first addition.  Since
a 32-bit add only ever touches 4 byte-rows, the table is materialised *on
demand*: each visited row is composed, decoded and written into storage once,
then cached.  The worst case is still complete; the common case is 4 rows.

**Exact keys, similarity fallback.**  A gate key is an exact XOR of
slot-and-position atoms, so it is looked up in a dictionary — a 131 072-row
similarity scan would cost more than the 32-stage ripple it replaces.  The
fuzzy path is still there (:meth:`SliceGate.resolve_similar`) because a
damaged key should degrade to "the nearest row" rather than fail; it is just
not on the hot path.  Any row resolution therefore reports a confidence, and
an exact hit reports 1.0.

Why the key does not collapse
-----------------------------
:mod:`hdc.adder` documents the parity bug: XORing a *shared* zero/one pair
loses everything but parity, so ``1^1^1`` and ``1^0^0`` collide.  Here every
(slot, position, value) triple has its own atom, so the key is an XOR of
distinct vectors encoding all ``2k+1`` inputs, and the 2**(2k+1) keys stay
distinct.
"""

from __future__ import annotations

from dataclasses import dataclass

from .hypervec import (
    CODEC_THRESHOLD,
    DEFAULT_DIM,
    bind,
    bundle,
    hvc,
    role,
    sim,
)
from .memory import HDCStorage

#: Above this many rows a gate is built lazily instead of eagerly.
EAGER_MAX_ROWS = 4096


@dataclass(frozen=True)
class SliceStep:
    """One multi-bit slice of the carry chain."""

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

    def render(self) -> str:
        mask = (1 << self.radix) - 1
        return (
            "slice %2d..%-2d a=0x%0*x b=0x%0*x cin=%d -> sum=0x%0*x cout=%d"
            "  (lookup sim %.3f)"
            % (
                self.index,
                self.index + self.radix - 1,
                self.radix // 4 or 1,
                self.a & mask,
                self.radix // 4 or 1,
                self.b & mask,
                self.carry_in,
                self.radix // 4 or 1,
                self.total & mask,
                self.carry_out,
                self.confidence,
            )
        )


@dataclass(frozen=True)
class SliceRow:
    """One materialised gate row."""

    a: int
    b: int
    carry_in: int
    total: int
    carry_out: int
    key: int
    cell: str

    def render(self) -> str:
        return "a=0x%x b=0x%x cin=%d -> sum=0x%x cout=%d  %s" % (
            self.a, self.b, self.carry_in, self.total, self.carry_out, self.cell
        )


class SliceGate:
    """A ``k``-bit adder slice whose truth table lives in HDC storage."""

    def __init__(
        self,
        radix: int = 8,
        d: int = DEFAULT_DIM,
        storage: HDCStorage | None = None,
        threshold: float = CODEC_THRESHOLD,
        eager: bool | None = None,
    ):
        if not 1 <= radix <= 8:
            raise ValueError("radix must be between 1 and 8 bits")
        self.radix = radix
        self.d = d
        self.threshold = threshold
        self.row_count = 1 << (2 * radix + 1)
        self.eager = (self.row_count <= EAGER_MAX_ROWS) if eager is None else eager
        self.storage = storage if storage is not None else HDCStorage(
            d=d, capacity=max(64, self.row_count if self.eager else 4096),
            label="gate%d" % radix,
        )
        #: One atom per (slot, position, value): this is what keeps keys distinct.
        self.a_atoms = {(i, b): hvc("gate/s%d/a/%d/%d" % (radix, i, b), d)
                        for i in range(radix) for b in (0, 1)}
        self.b_atoms = {(i, b): hvc("gate/s%d/b/%d/%d" % (radix, i, b), d)
                        for i in range(radix) for b in (0, 1)}
        self.cin_atoms = {b: hvc("gate/s%d/cin/%d" % (radix, b), d) for b in (0, 1)}
        self.cout_atoms = {b: hvc("gate/s%d/cout/%d" % (radix, b), d) for b in (0, 1)}
        self.sum_atoms = {(i, b): hvc("gate/s%d/sum/%d/%d" % (radix, i, b), d)
                          for i in range(radix) for b in (0, 1)}
        self.rows: dict[int, SliceRow] = {}
        self.lookups = 0
        self.materialised = 0
        self.similarity_fallbacks = 0
        if self.eager:
            self._build_all()

    # ------------------------------------------------------------- the gate
    def key_of(self, a: int, b: int, c: int) -> int:
        """Exact gate address: XOR of every slot/position/value atom."""
        key = 0
        for i in range(self.radix):
            key ^= self.a_atoms[(i, (a >> i) & 1)]
        for i in range(self.radix):
            key ^= self.b_atoms[(i, (b >> i) & 1)]
        key ^= self.cin_atoms[1 if c else 0]
        return key & ((1 << self.d) - 1)

    def _cell_name(self, a: int, b: int, c: int) -> str:
        width = (2 * self.radix + 1 + 3) // 4
        return "adder/s%d/%0*x" % (self.radix, width, (c << (2 * self.radix)) | (b << self.radix) | a)

    def _build_row(self, a: int, b: int, c: int) -> SliceRow:
        """Compose the row's record, write it to storage, and memoise it."""
        mask = (1 << self.radix) - 1
        total = (a + b + c) & mask
        carry = (a + b + c) >> self.radix
        key = self.key_of(a, b, c)
        fields = [
            bind(role("key", self.d), key, self.d),
            bind(role("carry", self.d), self.cout_atoms[carry], self.d),
        ]
        for i in range(self.radix):
            fields.append(
                bind(role("sum%d" % i, self.d), self.sum_atoms[(i, (total >> i) & 1)], self.d)
            )
        record = bundle(fields, self.d)
        cell = self._cell_name(a, b, c)
        self.storage.write(cell, record, symbol=True)
        row = SliceRow(a, b, c, total, carry, key, cell)
        self.rows[key] = row
        self.materialised += 1
        return row

    def _build_all(self) -> None:
        mask = (1 << self.radix) - 1
        for c in (0, 1):
            for b in range(mask + 1):
                for a in range(mask + 1):
                    self._build_row(a, b, c)

    def resolve(self, a: int, b: int, c: int) -> tuple[SliceRow, float]:
        """Look up the row for these inputs; materialise it on first use."""
        self.lookups += 1
        key = self.key_of(a, b, c)
        row = self.rows.get(key)
        if row is not None:
            return row, 1.0
        return self._build_row(a, b, c), 1.0

    def resolve_similar(self, a: int, b: int, c: int) -> tuple[SliceRow, float]:
        """Content-addressable fallback: nearest materialised row.

        Only useful when a key arrives damaged (or comes from outside this
        gate).  Exact keys resolve at 1.0 through :meth:`resolve`; this path
        reports the similarity it actually achieved.
        """
        key = self.key_of(a, b, c)
        if key in self.rows:
            return self.rows[key], 1.0
        best_row, best = None, -2.0
        for candidate, row in self.rows.items():
            s = sim(candidate, key, self.d)
            if s > best:
                best, best_row = s, row
        if best_row is None:
            self.similarity_fallbacks += 1
            return self._build_row(a, b, c), 0.0
        self.similarity_fallbacks += 1
        return best_row, best

    # ------------------------------------------------------------- evaluate
    def evaluate(self, a: int, b: int, c: int) -> tuple[int, int, float]:
        """Run one slice: resolve the row, then read its fields back out."""
        row, key_confidence = self.resolve(a, b, c)
        record = self.storage.read(row.cell)
        carry, carry_margin = self._read_bit(
            bind(record, role("carry", self.d), self.d), self.cout_atoms
        )
        total = 0
        margins = [carry_margin]
        for i in range(self.radix):
            bit, margin = self._read_bit(
                bind(record, role("sum%d" % i, self.d), self.d),
                {0: self.sum_atoms[(i, 0)], 1: self.sum_atoms[(i, 1)]},
            )
            total |= bit << i
            margins.append(margin)
        if (row.total, row.carry_out) != (total, carry):
            raise AssertionError(
                "gate row read back as %s but was built as %s"
                % ((total, carry), (row.total, row.carry_out))
            )
        return total, carry, min(min(margins), key_confidence)

    def _read_bit(self, probe: int, atoms: dict) -> tuple[int, float]:
        s0 = sim(probe, atoms[0], self.d)
        s1 = sim(probe, atoms[1], self.d)
        if s1 >= s0:
            return 1, s1 - s0
        return 0, s0 - s1

    # ---------------------------------------------------------------- views
    def summary(self) -> dict:
        return {
            "radix": self.radix,
            "stages_for_32": (32 + self.radix - 1) // self.radix,
            "table_rows": self.row_count,
            "materialised": len(self.rows),
            "eager": self.eager,
            "lookups": self.lookups,
            "similarity_fallbacks": self.similarity_fallbacks,
            "cells": len(self.storage),
        }

    def table_text(self, limit: int = 12) -> str:
        head = "gate radix=%d: %d rows total, %d materialised%s" % (
            self.radix,
            self.row_count,
            len(self.rows),
            " (eager)" if self.eager else " (lazy)",
        )
        rows = [head]
        for row in list(self.rows.values())[:limit]:
            rows.append("  " + row.render())
        if len(self.rows) > limit:
            rows.append("  ... %d more rows" % (len(self.rows) - limit))
        return "\n".join(rows)