"""Tests for the read-back kernel ``Phi`` (``palimseste.phi``).

Key behaviors:
  - Phi returns the stored value when queried with the exact stored address
  - Phi blends multiple matching values via weighted bundle (majority)
  - Phi returns None on empty memory / no matches (no fabrication)
  - radius controls neighborhood width
  - min_weight + topk cutoffs work
  - reconstruct_param implements Axiome 4 (W_t = Phi(bind(x, s_t)))
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste import hv
from palimseste.memory import Memory
from palimseste.phi import Phi, KernelConfig, Retrieval


def _mem_with(D=2000, n=0, seed=0) -> tuple[Memory, np.random.Generator]:
    rng = np.random.default_rng(seed)
    return Memory(D=D, rng=np.random.default_rng(seed)), rng


def test_phi_exact_recall():
    # query the exact stored address -> recover its value
    mem, rng = _mem_with(D=2000, seed=1)
    addr = hv.random_hv(D=2000, rng=rng)
    val = hv.random_hv(D=2000, rng=rng)
    mem.write(addr, val)
    phi = Phi(config=KernelConfig(radius=0))
    out = phi(mem, addr)
    assert out is not None
    # exact match at radius 0 -> identical value
    assert out == val


def test_phi_empty_memory_returns_none():
    mem, _ = _mem_with(D=500, seed=2)
    phi = Phi(config=KernelConfig(radius=50))
    assert phi(mem, hv.random_hv(D=500)) is None


def test_phi_no_match_returns_none():
    mem, rng = _mem_with(D=2000, seed=3)
    mem.write(hv.random_hv(D=2000, rng=rng), hv.random_hv(D=2000, rng=rng))
    phi = Phi(config=KernelConfig(radius=0))
    # query a fresh random HV -> almost surely > 0 Hamming away
    q = hv.random_hv(D=2000, rng=rng)
    # try a few times; with radius 0 and random q, no match is expected
    out = phi(mem, q)
    # could rarely collide; assert None or equality with the single value
    assert out is None or out == mem.traces[0].value


def test_phi_blends_multiple_values():
    # Three traces near the same address with different values:
    # majority value should win the bundle.
    mem, rng = _mem_with(D=3000, seed=4)
    base_addr = hv.random_hv(D=3000, rng=rng)
    # create near addresses by flipping a few bits
    def near(h, n):
        s = hv.bits_to_signs(h)
        pos = rng.choice(h.D, size=n, replace=False)
        s[pos] = -s[pos]
        return hv.signs_to_bits(s)

    v_majority = hv.random_hv(D=3000, rng=rng)
    v_minority = hv.random_hv(D=3000, rng=rng)
    # two traces point to v_majority, one to v_minority
    mem.write(near(base_addr, 3), v_majority)
    mem.write(near(base_addr, 4), v_majority)
    mem.write(near(base_addr, 5), v_minority)
    phi = Phi(config=KernelConfig(radius=20))
    out = phi(mem, base_addr)
    assert out is not None
    # majority should dominate
    assert hv.similarity(out, v_majority) > hv.similarity(out, v_minority)


def test_phi_radius_filters():
    mem, rng = _mem_with(D=2000, seed=5)
    addr = hv.random_hv(D=2000, rng=rng)
    val = hv.random_hv(D=2000, rng=rng)
    mem.write(addr, val)
    # query with something ~10% different; radius 0 excludes it, radius large includes
    s = hv.bits_to_signs(addr)
    pos = rng.choice(addr.D, size=200, replace=False)
    s[pos] = -s[pos]
    q = hv.signs_to_bits(s)
    phi0 = Phi(config=KernelConfig(radius=0))
    phi_big = Phi(config=KernelConfig(radius=500))
    assert phi0(mem, q) is None
    out = phi_big(mem, q)
    assert out is not None


def test_phi_topk_limits_matches():
    mem, rng = _mem_with(D=2000, seed=6)
    base = hv.random_hv(D=2000, rng=rng)

    def near(h, n):
        s = hv.bits_to_signs(h)
        pos = rng.choice(h.D, size=n, replace=False)
        s[pos] = -s[pos]
        return hv.signs_to_bits(s)

    for _ in range(20):
        mem.write(near(base, 5), hv.random_hv(D=2000, rng=rng))
    phi = Phi(config=KernelConfig(radius=50, topk=3))
    ret = phi.retrieve(mem, base)
    assert ret.n_matches <= 3


def test_phi_min_weight_filters():
    mem, rng = _mem_with(D=2000, seed=7)
    addr = hv.random_hv(D=2000, rng=rng)
    val = hv.random_hv(D=2000, rng=rng)
    mem.write(addr, val, weight=1e-5)  # below default min_weight 1e-3
    phi = Phi(config=KernelConfig(radius=0, min_weight=1e-3))
    assert phi(mem, addr) is None
    # lowering the floor recovers it
    phi_loose = Phi(config=KernelConfig(radius=0, min_weight=1e-7))
    assert phi_loose(mem, addr) == val


def test_phi_reconstruct_param_axiom4():
    # W_t = Phi(bind(x, s_t))
    mem, rng = _mem_with(D=2000, seed=8)
    x = hv.random_hv(D=2000, rng=rng)
    s = hv.random_hv(D=2000, rng=rng)
    W = hv.random_hv(D=2000, rng=rng)
    # store W under address bind(x, s)
    mem.write(hv.bind(x, s), W)
    phi = Phi(config=KernelConfig(radius=0))
    rec = phi.reconstruct_param(mem, x, s)
    assert rec is not None
    assert rec == W


def test_kernel_config_encode_decode_roundtrip():
    c = KernelConfig(radius=42, min_weight=0.01, sharpness=2.0, topk=5)
    d = c.encode()
    c2 = KernelConfig.decode(d)
    assert c2 == c


def test_kernel_config_invalid():
    with pytest.raises(ValueError):
        KernelConfig(radius=-1)
    with pytest.raises(ValueError):
        KernelConfig(min_weight=0.0)
    with pytest.raises(ValueError):
        KernelConfig(topk=0)


def test_retrieval_introspection():
    mem, rng = _mem_with(D=2000, seed=9)
    addr = hv.random_hv(D=2000, rng=rng)
    val = hv.random_hv(D=2000, rng=rng)
    mem.write(addr, val)
    phi = Phi(config=KernelConfig(radius=10))
    ret = phi.retrieve(mem, addr)
    assert isinstance(ret, Retrieval)
    assert ret.n_matches == 1
    assert ret.sims[0] == pytest.approx(1.0)
    assert ret.weights[0] == pytest.approx(1.0)