"""Tests for the Learner + Encoder (``palimseste.learner``).

Verifies:
  - learn() is O(1): the millionth example costs the same as the first
    (we measure write cost stays flat across |M| scales)
  - predict() recovers learned mappings
  - context binding: same x in different ctx yields different addresses
  - Encoder: ints/strs/floats/sequences/sets encode deterministically &
    with the right similarity structure (floats level-coded, sequences
    order-sensitive, sets order-insensitive)
"""

from __future__ import annotations

import time
import numpy as np
import pytest

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


def _learner(D=2000, seed=0, radius=0) -> tuple[Learner, np.random.Generator]:
    rng = np.random.default_rng(seed)
    mem = Memory(D=D, rng=np.random.default_rng(seed))
    phi = Phi(config=KernelConfig(radius=radius, min_weight=1e-6))
    return Learner(mem=mem, phi=phi, rng=rng), rng


# ----------------------------------------------------------------- learner
def test_learn_then_predict_exact():
    lr, rng = _learner(D=2000, radius=0)
    x = hv.random_hv(D=2000, rng=rng)
    y = hv.random_hv(D=2000, rng=rng)
    lr.learn(x, y)
    p = lr.predict(x)
    assert p.value == y
    assert p.n_matches == 1
    assert p.confidence > 0.99


def test_learn_context_binding():
    # same x, different ctx -> different addresses, different predictions
    lr, rng = _learner(D=2000, radius=0)
    x = hv.random_hv(D=2000, rng=rng)
    ctx_a = hv.random_hv(D=2000, rng=rng)
    ctx_b = hv.random_hv(D=2000, rng=rng)
    y_a = hv.random_hv(D=2000, rng=rng)
    y_b = hv.random_hv(D=2000, rng=rng)
    lr.learn(x, y_a, ctx=ctx_a)
    lr.learn(x, y_b, ctx=ctx_b)
    pa = lr.predict(x, ctx=ctx_a)
    pb = lr.predict(x, ctx=ctx_b)
    assert pa.value == y_a
    assert pb.value == y_b


def test_learn_is_o1_flat_cost():
    # The write cost must not grow with |M|. We compare per-write time at
    # small |M| vs large |M| and require they are within a generous factor.
    D = 4000
    rng = np.random.default_rng(7)
    mem = Memory(D=D, rng=np.random.default_rng(7))
    phi = Phi(config=KernelConfig(radius=10))
    lr = Learner(mem=mem, phi=phi, rng=rng)

    xs = [hv.random_hv(D=D, rng=rng) for _ in range(2000)]
    ys = [hv.random_hv(D=D, rng=rng) for _ in range(2000)]

    t0 = time.perf_counter()
    for i in range(50):
        lr.learn(xs[i], ys[i])
    small = (time.perf_counter() - t0) / 50

    for i in range(50, 2000):
        lr.learn(xs[i], ys[i])

    t1 = time.perf_counter()
    for i in range(2000, 2050):
        lr.learn(xs[i % len(xs)], ys[i % len(ys)])
    large = (time.perf_counter() - t1) / 50

    # O(1) amortized: large-|M| writes should be within 5x of small-|M| writes
    # (generous to absorb LSH dict growth / GC noise).
    ratio = large / max(small, 1e-9)
    assert ratio < 5.0, f"write cost grew: small={small:.2e} large={large:.2e} ratio={ratio:.2f}"


def test_learn_many_appends_all():
    lr, rng = _learner(D=1000, radius=5)
    pairs = [(hv.random_hv(D=1000, rng=rng), hv.random_hv(D=1000, rng=rng)) for _ in range(30)]
    traces = lr.learn_many(pairs)
    assert len(traces) == 30
    assert len(lr.mem) == 30


def test_predict_empty_returns_none():
    lr, _ = _learner(D=500)
    p = lr.predict(hv.random_hv(D=500))
    assert p.value is None
    assert p.confidence == 0.0
    assert p.n_matches == 0


# ----------------------------------------------------------------- encoder
def _enc(D=2000, seed=3) -> Encoder:
    return Encoder(D=D, rng=np.random.default_rng(seed))


def test_encode_int_deterministic():
    e = _enc()
    assert e.encode_int(5) == e.encode_int(5)
    assert e.encode_int(5) != e.encode_int(6)


def test_encode_str_deterministic():
    e = _enc()
    assert e.encode_str("cat") == e.encode_str("cat")
    assert e.encode_str("cat") != e.encode_str("dog")


def test_encode_float_level_similarity():
    # adjacent floats are similar, distant are not.
    # With 256 levels and ~0.025*D bits/step, the random-walk level code
    # saturates to quasi-orthogonal within ~30 levels, so values far apart
    # in [0,1] decorrelate as VSA scalar coding requires.
    e = _enc(D=5000, seed=1)
    a = e.encode_float(0.41)
    b = e.encode_float(0.42)
    c = e.encode_float(0.90)
    assert hv.similarity(a, b) > 0.5
    # 0.41 vs 0.90 are ~126 levels apart -> well past saturation
    assert hv.similarity(a, c) < 0.3


def test_encode_float_clamps():
    e = _enc()
    # out-of-range clamps, no error
    assert isinstance(e.encode_float(-5.0), hv.HV)
    assert isinstance(e.encode_float(99.0), hv.HV)


def test_encode_float_bad_range():
    e = _enc()
    with pytest.raises(ValueError):
        e.encode_float(0.5, lo=1.0, hi=0.0)


def test_encode_sequence_order_sensitive():
    e = _enc(D=5000, seed=2)
    a = e.encode_int(1)
    b = e.encode_int(2)
    c = e.encode_int(3)
    seq_abc = e.encode_sequence([a, b, c])
    seq_cba = e.encode_sequence([c, b, a])
    # order matters: different permutations are dissimilar
    assert hv.similarity(seq_abc, seq_cba) < 0.6


def test_encode_set_order_insensitive():
    e = _enc(D=5000, seed=3)
    a = e.encode_int(1)
    b = e.encode_int(2)
    c = e.encode_int(3)
    s1 = e.encode_set([a, b, c])
    s2 = e.encode_set([c, b, a])
    # same set, different order -> similar (not necessarily equal due to tie-breaks,
    # but strongly correlated)
    assert hv.similarity(s1, s2) > 0.7


def test_encode_kv():
    # Bundling is commutative ONLY when there are no ties. With an even
    # number of independent bound-pairs every coordinate is an exact tie and
    # the result degenerates to random tie-breaking (sim ~ 0 between orderings).
    # With an odd count (>=3) ties are broken by genuine majority and the
    # bundle is order-insensitive.
    e = _enc(D=5000, seed=4)
    pairs = [
        (e.encode_str("color"), e.encode_str("red")),
        (e.encode_str("size"), e.encode_int(42)),
        (e.encode_str("weight"), e.encode_int(7)),
        (e.encode_str("shape"), e.encode_str("cube")),
        (e.encode_str("owner"), e.encode_str("ada")),
    ]
    rec = e.encode_kv(pairs)
    rec2 = e.encode_kv(list(reversed(pairs)))
    assert hv.similarity(rec, rec2) > 0.7


def test_prediction_dataclass():
    p = Prediction(query=hv.random_hv(10), value=None, confidence=0.0, n_matches=0)
    assert p.value is None
    assert p.confidence == 0.0