"""PALIMPSESTE — Capacity & scaling benchmarks.

Verifies the spec's core performance claims:
  1. O(1) write cost: the millionth insertion costs the same as the first.
  2. Sub-linear retrieval: Phi query time grows slowly with |M| (LSH).
  3. Capacity: recall accuracy stays high as |M| grows into the thousands
     (well below the exponential theoretical ceiling, but demonstrating the
     "we don't saturate" property at laptop scale).

Run:  python examples/benchmark.py
"""

from __future__ import annotations

import time
import numpy as np

from palimseste import hv
from palimseste.memory import Memory
from palimseste.phi import Phi, KernelConfig
from palimseste.learner import Learner


def bench_o1_write(D: int = 4000, sizes=(100, 500, 1000, 2000)) -> None:
    print("\n--- O(1) write cost vs |M| ---")
    print(f"{'|M|':>8} {'per-write (us)':>16} {'ratio vs first':>16}")
    rng = np.random.default_rng(0)
    mem = Memory(D=D, rng=np.random.default_rng(0))
    phi = Phi(config=KernelConfig(radius=8, min_weight=1e-6))
    lr = Learner(mem=mem, phi=phi, rng=rng)
    xs = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    ys = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    base = None
    i = 0
    for target in sizes:
        while len(mem) < target:
            lr.learn(xs[i], ys[i])
            i += 1
        # time 50 writes
        t0 = time.perf_counter()
        for _ in range(50):
            lr.learn(xs[i % len(xs)], ys[i % len(ys)])
            i += 1
        dt = (time.perf_counter() - t0) / 50 * 1e6
        if base is None:
            base = dt
        print(f"{len(mem):>8} {dt:>16.1f} {dt/base:>16.2f}")


def bench_retrieval(D: int = 4000, sizes=(100, 500, 1000, 2000)) -> None:
    print("\n--- Phi retrieval time vs |M| ---")
    print(f"{'|M|':>8} {'per-query (ms)':>16} {'cand/|M|':>10}")
    rng = np.random.default_rng(1)
    mem = Memory(D=D, rng=np.random.default_rng(1))
    phi = Phi(config=KernelConfig(radius=20, min_weight=1e-6))
    addrs = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    vals = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    i = 0
    for target in sizes:
        while len(mem) < target:
            mem.write(addrs[i], vals[i])
            i += 1
        q = addrs[i % len(addrs)]
        # warm
        phi(mem, q)
        t0 = time.perf_counter()
        for _ in range(50):
            phi(mem, q)
        dt = (time.perf_counter() - t0) / 50 * 1e3
        cand = mem.candidates(q)
        print(f"{len(mem):>8} {dt:>16.3f} {len(cand)/len(mem):>10.3f}")


def bench_capacity(D: int = 4000, sizes=(100, 500, 1000, 2000), radius=0) -> None:
    print("\n--- Recall accuracy vs |M| (exact-address recall) ---")
    print(f"{'|M|':>8} {'recall@1':>10} {'recall@r':>10}")
    rng = np.random.default_rng(2)
    mem = Memory(D=D, rng=np.random.default_rng(2))
    phi_exact = Phi(config=KernelConfig(radius=0, min_weight=1e-6))
    phi_wide = Phi(config=KernelConfig(radius=radius if radius > 0 else 30, min_weight=1e-6))
    addrs = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    vals = [hv.random_hv(D=D, rng=rng) for _ in range(max(sizes))]
    i = 0
    for target in sizes:
        while len(mem) < target:
            mem.write(addrs[i], vals[i])
            i += 1
        # sample 100 stored addresses and check exact recall
        idx = rng.choice(len(mem), size=min(100, len(mem)), replace=False)
        ok_exact = ok_wide = 0
        for k in idx:
            q = mem.traces[k].address
            out_e = phi_exact(mem, q)
            out_w = phi_wide(mem, q)
            if out_e == mem.traces[k].value:
                ok_exact += 1
            if out_w == mem.traces[k].value:
                ok_wide += 1
        n = len(idx)
        print(f"{len(mem):>8} {ok_exact/n:>10.3f} {ok_wide/n:>10.3f}")


def main() -> None:
    print("=" * 64)
    print("PALIMPSESTE — capacity & scaling benchmarks")
    print("=" * 64)
    bench_o1_write()
    bench_retrieval()
    bench_capacity()
    print("\n" + "=" * 64)
    print("Done. Theoretical capacity (Kanerva SDM): C ~ 0.14*D * 2^(D/beta).")
    print("With D=10000, this vastly exceeds any laptop-scale |M| tested here.")
    print("=" * 64)


if __name__ == "__main__":
    main()