import mlx.core as mx

from morph_hrr import (
    bind,
    bundle,
    cosine_similarity,
    make_unitary,
    normalize,
    unbind,
    update_context,
)

DIM = 4096


def _close(actual, expected, tol):
    assert abs(float(actual) - float(expected)) <= tol, f"{actual} !~= {expected} (tol {tol})"


def test_unitary_keys_have_unit_fft_magnitude():
    key = make_unitary(DIM, seed=10)
    mags = mx.abs(mx.fft.fft(key))
    _close(mx.min(mags), 1.0, 1e-4)
    _close(mx.max(mags), 1.0, 1e-4)


def test_unitary_bind_unbind_recovers_value():
    key = make_unitary(DIM, seed=11)
    value = normalize(mx.random.normal((DIM,)))
    recovered = unbind(bind(key, value), key)
    assert float(cosine_similarity(value, recovered)) > 0.995


def test_bind_unbind_supports_batch():
    key = make_unitary(DIM, seed=12)
    values = normalize(mx.random.normal((3, DIM)))
    recovered = unbind(bind(key, values), key)
    assert recovered.shape == values.shape
    assert float(mx.min(cosine_similarity(values, recovered))) > 0.995


def test_bundle_recovery_beats_distractors():
    roles = [make_unitary(DIM, seed=13 + i) for i in range(3)]
    vals = [normalize(mx.random.normal((DIM,))) for _ in range(3)]
    memory = bundle(*(bind(r, v) for r, v in zip(roles, vals)))
    recovered = unbind(memory, roles[1])
    assert float(cosine_similarity(vals[1], recovered)) > 0.50
    assert abs(float(cosine_similarity(vals[0], recovered))) < 0.10
    assert abs(float(cosine_similarity(vals[2], recovered))) < 0.10


def test_normalize_yields_unit_norm():
    v = normalize(mx.random.normal((DIM,)))
    _close(mx.sum(v * v), 1.0, 1e-4)


def test_update_context_stays_normalized_and_finite():
    context = mx.zeros((DIM,), dtype=mx.float32)
    for i in range(128):
        token = normalize(mx.random.normal((DIM,)))
        context = update_context(context, token, make_unitary(DIM, seed=1_000 + i))
    _close(mx.sum(context * context), 1.0, 1e-4)
    assert bool(mx.all(mx.isfinite(context)))


def test_empty_bundle_raises():
    try:
        bundle()
    except ValueError:
        return
    raise AssertionError("empty bundle must raise ValueError")


def test_determinism_same_seed():
    a = make_unitary(DIM, seed=42)
    b = make_unitary(DIM, seed=42)
    assert bool(mx.allclose(a, b))