"""Photographic, holographic and hybrid memory.

Three ways to remember the same thing, with the trade stated as a measurement
rather than a slogan.

**Photographic** — one cell per item, exact.  Recall returns the stored vector
bit for bit.  Cost: ``O(1)`` storage per item and a hard capacity.  This is the
right default: in this machine a cell already holds a full 10 000-bit vector, so
degrading recall to save space buys nothing.

**Holographic** — every item is bound into one superposition and the trace is
recovered by *majority*, per bit:

    trace_bit = 1 if more than half of the bound terms have that bit set

Unbinding one key leaves the item plus crosstalk from the others, so recall is
a cleanup step: argmax similarity over a codebook.  Majority is the whole
point.  The obvious alternative, XOR accumulation, does **not** work: with
``trace ^= bind(k, v)`` the crosstalk is an XOR of random terms, which is
itself random at every bit, so recall is chance at any k — measured at 0%
below.  Majority crosstalk instead grows as ``sqrt(k)`` against a signal of
``d``, so similarity decays smoothly and retrieval keeps working for a while.

**Hybrid** — photographic buckets, each with its own holographic trace.  The
exact index answers hits at full fidelity; the trace answers approximate
queries and survives index eviction.  That is a cache with a degradation path
rather than a cliff.
"""

from __future__ import annotations

from typing import Iterable, Sequence

from .hypervec import DEFAULT_DIM, bind, hvc, sim


def _bits(value: int, d: int) -> int:
    return value & ((1 << d) - 1)


class PhotographicMemory:
    """Exact, one cell per key. Recall is bit-identical or a miss."""

    def __init__(self, d: int = DEFAULT_DIM, capacity: int | None = None):
        self.d = d
        self.capacity = capacity
        self.cells: dict[str, int] = {}
        self.evicted = 0

    def capture(self, key: str, value: int) -> None:
        """Store a value exactly. Oldest key is evicted when full."""
        if self.capacity is not None and key not in self.cells and len(self.cells) >= self.capacity:
            oldest = next(iter(self.cells))
            del self.cells[oldest]
            self.evicted += 1
        self.cells[key] = _bits(value, self.d)

    def recall(self, key: str) -> tuple[int | None, float]:
        """``(value, 1.0)`` on a hit, ``(None, 0.0)`` on a miss."""
        value = self.cells.get(key)
        return (value, 1.0) if value is not None else (None, 0.0)

    def __contains__(self, key: str) -> bool:
        return key in self.cells

    def __len__(self) -> int:
        return len(self.cells)

    def stats(self) -> dict:
        return {"mode": "photographic", "items": len(self.cells), "capacity": self.capacity,
                "evicted": self.evicted, "exact": True}


class HolographicMemory:
    """One superposition for everything, held as per-bit majority counts."""

    def __init__(self, d: int = DEFAULT_DIM, alphabet: Sequence[int] | None = None):
        self.d = d
        self.alphabet = [_bits(v, d) for v in (alphabet or [])]
        # bit -> number of bound terms with that bit set; the trace is the
        # majority of these, so it is O(d) state no matter how many items.
        self._ones = [0] * d
        #: key -> value for every key currently in the trace, so re-recording
        #: a key *replaces* it instead of piling up votes.  Without this, a
        #: repeated observation outweighs every other pair and the majority
        #: stops being a superposition of distinct keys at all.
        self._bound: dict[str, int] = {}
        self.keys: list[str] = []
        #: raw key vectors, so a graded query can snap to the nearest one
        self.key_vectors: list[int] = []

    @property
    def items(self) -> int:
        """Distinct keys in the trace."""
        return len(self._bound)

    def _keyvector(self, key: str | int) -> int:
        """A key is either a symbol (hashed) or a vector used as-is.

        Using the vector as its own key is what lets *similar* queries resolve
        to a stored entry: hashing a state to a string throws that similarity
        away, which is fatal for graded inputs (contexts, images, embeddings).
        """
        return key if isinstance(key, int) else hvc(key, self.d)

    def record(self, key: str | int, value: int) -> None:
        """Bind one pair into the trace, replacing any earlier value for ``key``."""
        bits = _bits(value, self.d)
        ones = self._ones
        old = self._bound.get(key)
        if old is not None:
            for i in range(self.d):
                if (old >> i) & 1:
                    ones[i] -= 1
        bound = bind(self._keyvector(key), bits, self.d)
        for i in range(self.d):
            if (bound >> i) & 1:
                ones[i] += 1
        if old is None:
            self.keys.append(key)
            self.key_vectors.append(self._keyvector(key))
        self._bound[key] = bits

    @property
    def trace(self) -> int:
        """The majority vector: the superposition itself."""
        half = len(self._bound)
        out = 0
        for i in range(self.d):
            if self._ones[i] * 2 > half:
                out |= 1 << i
        return out

    def _unbind(self, key: str | int) -> int:
        return bind(self.trace, self._keyvector(key), self.d)

    def _cleanup(self, probe: int) -> tuple[int | None, float]:
        best, best_sim = None, -2.0
        for candidate in self.alphabet:
            s = sim(probe, candidate, self.d)
            if s > best_sim:
                best, best_sim = candidate, s
        return best, best_sim

    def recall(self, key: str | int) -> tuple[int | None, float]:
        """Unbind, then clean up against the codebook."""
        if not self.keys:
            return None, 0.0
        probe = self._unbind(key)
        if not self.alphabet:
            return probe, sim(probe, self.trace, self.d)
        return self._cleanup(probe)

    def nearest_key(self, query: int) -> tuple[int | None, float]:
        """The stored key vector closest to ``query`` and its similarity."""
        best, best_sim = None, -2.0
        for vector in self.key_vectors:
            score = sim(query, vector, self.d)
            if score > best_sim:
                best, best_sim = vector, score
        return best, best_sim

    def recall_bound(self, key_vector: int) -> tuple[int | None, float]:
        """Unbind with a key *vector* rather than a symbol, then clean up."""
        if not self.keys:
            return None, 0.0
        probe = bind(self.trace, key_vector, self.d)
        if not self.alphabet:
            return probe, sim(probe, self.trace, self.d)
        return self._cleanup(probe)

    def confidence(self, key: str | int) -> float:
        """Cleanup similarity for ``key``: how well it matches its best symbol.

        This is the presence signal.  Comparing the raw unbinding to the trace
        does *not* work for XOR-style binding: the two are unrelated vectors.
        """
        if not self.keys:
            return 0.0
        probe = self._unbind(key)
        if self.alphabet:
            return self._cleanup(probe)[1]
        return sim(probe, self.trace, self.d)

    def stats(self) -> dict:
        return {"mode": "holographic", "items": self.items,
                "codebook": len(self.alphabet), "exact": False}


class HybridMemory:
    """Photographic buckets, each backed by its own holographic trace."""

    def __init__(self, d: int = DEFAULT_DIM, buckets: int = 8, per_bucket: int | None = None,
                 alphabet: Sequence[int] | None = None):
        if buckets < 1:
            raise ValueError("need at least one bucket")
        self.d = d
        self.nbuckets = buckets
        self.per_bucket = per_bucket
        self.alphabet = list(alphabet or [])
        self.exact: list[dict[str, int]] = [{} for _ in range(buckets)]
        self.traces: list[HolographicMemory] = [
            HolographicMemory(d=d, alphabet=self.alphabet) for _ in range(buckets)
        ]
        self.fallbacks = 0

    def _bucket(self, key: str) -> int:
        digest = 0
        for ch in key:
            digest = (digest * 131 + ord(ch)) & 0xFFFFFFFF
        return digest % self.nbuckets

    def capture(self, key: str, value: int) -> None:
        b = self._bucket(key)
        cell = self.exact[b]
        if self.per_bucket is not None and key not in cell and len(cell) >= self.per_bucket:
            oldest = next(iter(cell))
            del cell[oldest]
            self._rebuild(b)
        cell[key] = _bits(value, self.d)
        if cell is self.exact[b] and key in cell:
            self.traces[b].record(key, cell[key])

    def _rebuild(self, b: int) -> None:
        """Rebuild one trace from its surviving exact entries."""
        self.traces[b] = HolographicMemory(d=self.d, alphabet=self.alphabet)
        for key, value in self.exact[b].items():
            self.traces[b].record(key, value)

    def recall(self, key: str) -> tuple[int | None, float, str]:
        """``(value, similarity, path)`` where path is ``exact``, ``trace`` or ``miss``."""
        b = self._bucket(key)
        value = self.exact[b].get(key)
        if value is not None:
            return value, 1.0, "exact"
        trace = self.traces[b]
        if not trace.keys:
            return None, 0.0, "miss"
        self.fallbacks += 1
        best, score = trace.recall(key)
        return best, score, "trace" if best is not None else "miss"

    def consolidate(self) -> int:
        """Rebuild every trace from the exact index; returns traces rewritten."""
        rebuilt = 0
        for b in range(self.nbuckets):
            before = self.traces[b]._ones
            self._rebuild(b)
            if self.traces[b]._ones != before:
                rebuilt += 1
        return rebuilt

    def stats(self) -> dict:
        exact = sum(len(cell) for cell in self.exact)
        return {"mode": "hybrid", "buckets": self.nbuckets, "exact_items": exact,
                "traced_buckets": sum(1 for t in self.traces if t.keys),
                "fallbacks": self.fallbacks, "exact": False}


# --------------------------------------------------------------------------
# measurement

def capacity_report(d: int = 2000, sizes: Iterable[int] = (2, 4, 8, 16, 32), symbols: int = 40) -> list[dict]:
    """Measure holographic cleanup accuracy against items sharing one trace."""
    alphabet = [hvc("codebook/%d" % i, d) for i in range(symbols)]
    rows = []
    for k in sizes:
        mem = HolographicMemory(d=d, alphabet=alphabet)
        pairs = [("k/%d" % i, alphabet[i % symbols]) for i in range(k)]
        for key, value in pairs:
            mem.record(key, value)
        hits = 0
        scores = []
        for key, value in pairs:
            got, score = mem.recall(key)
            scores.append(score)
            if got == value:
                hits += 1
        rows.append({"items": k, "exact_recall": hits / len(pairs),
                     "mean_similarity": sum(scores) / len(scores), "codebook": symbols})
    return rows


def compare_modes(d: int = 2000, k: int = 8, symbols: int = 8) -> dict:
    """Same items into all three modes, measured side by side."""
    alphabet = [hvc("symbol/%d" % i, d) for i in range(symbols)]
    keys = ["item/%d" % i for i in range(k)]
    values = [alphabet[i % symbols] for i in range(k)]

    photo = PhotographicMemory(d=d)
    holo = HolographicMemory(d=d, alphabet=alphabet)
    hybrid = HybridMemory(d=d, buckets=4, alphabet=alphabet)
    for key, value in zip(keys, values):
        photo.capture(key, value)
        holo.record(key, value)
        hybrid.capture(key, value)

    def score(mem, key, value, mode):
        if mode == "hybrid":
            got, _sim, _path = mem.recall(key)
        else:
            got, _sim = mem.recall(key)
        return got == value

    return {
        "items": k,
        "photographic": sum(score(photo, k_, v, "photo") for k_, v in zip(keys, values)) / k,
        "holographic": sum(score(holo, k_, v, "holo") for k_, v in zip(keys, values)) / k,
        "hybrid": sum(score(hybrid, k_, v, "hybrid") for k_, v in zip(keys, values)) / k,
    }


def demo_report(d: int = 2000) -> str:
    """Printable comparison (used by ``python -m hdc memory``)."""
    lines = ["memory modes at d=%d" % d]

    photo = PhotographicMemory(d=d)
    for i in range(4):
        photo.capture("frame/%d" % i, hvc("pixel/%d" % i, d))
    value, score = photo.recall("frame/2")
    lines.append("photographic: hit=%s similarity=%.2f items=%d"
                 % (value is not None, score, len(photo)))

    alphabet = [hvc("symbol/%d" % i, d) for i in range(8)]
    holo = HolographicMemory(d=d, alphabet=alphabet)
    for i in range(4):
        holo.record("slot/%d" % i, alphabet[i])
    hits = sum(1 for i in range(4) if holo.recall("slot/%d" % i)[0] == alphabet[i])
    lines.append("holographic: %d/4 exact over an 8-symbol codebook, confidence %.2f"
                 % (hits, holo.confidence("slot/1")))

    hybrid = HybridMemory(d=d, buckets=4, per_bucket=3, alphabet=alphabet)
    for i in range(6):
        hybrid.capture("key/%d" % i, alphabet[i])
    _v, score5, path5 = hybrid.recall("key/5")
    _v, score0, path0 = hybrid.recall("key/0")
    lines.append("hybrid: key/5 via %s (sim %.2f), key/0 via %s" % (path5, score5, path0))
    lines.append("hybrid consolidate rewrote %d traces" % hybrid.consolidate())

    lines.append("holographic capacity, 40-symbol codebook:")
    for row in capacity_report(d=d):
        lines.append("  k=%-3d exact %.0f%%  mean cleanup similarity %.3f"
                     % (row["items"], 100 * row["exact_recall"], row["mean_similarity"]))
    return "\n".join(lines)


if __name__ == "__main__":  # pragma: no cover - manual inspection
    print(demo_report())