"""Tests for hypervector primitives (``palimseste.hv``).

These encode the algebraic invariants the whole substrate relies on:
  - bind is self-inverse (reversible / addressable)
  - similarity of random HVs concentrates near 0 (quasi-orthogonality)
  - bundle is majority vote with random tie-breaks
  - packed-bit representation round-trips correctly
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste import hv


# ----------------------------------------------------------------- repr/shape
def test_HV_rejects_bad_length():
    with pytest.raises(ValueError):
        hv.HV(bits=np.zeros(5, dtype=np.uint8), D=100)  # 5 != ceil(100/8)=13


def test_HV_rejects_bad_dtype():
    with pytest.raises(TypeError):
        hv.HV(bits=np.zeros(13, dtype=np.int32), D=100)


def test_HV_rejects_2d():
    with pytest.raises(ValueError):
        hv.HV(bits=np.zeros((2, 6), dtype=np.uint8), D=96)


# ------------------------------------------------------------ construction
def test_random_hv_is_bipolar():
    rng = np.random.default_rng(0)
    x = hv.random_hv(D=1000, rng=rng)
    s = hv.bits_to_signs(x)
    assert set(np.unique(s).tolist()) <= {1, -1}
    assert x.D == 1000


def test_constant_hv():
    p = hv.constant_hv(D=64, value=+1)
    m = hv.constant_hv(D=64, value=-1)
    assert hv.bits_to_signs(p).tolist() == [1] * 64
    assert hv.bits_to_signs(m).tolist() == [-1] * 64
    with pytest.raises(ValueError):
        hv.constant_hv(D=64, value=0)


def test_bits_signs_roundtrip():
    rng = np.random.default_rng(1)
    x = hv.random_hv(D=500, rng=rng)
    signs = hv.bits_to_signs(x)
    x2 = hv.signs_to_bits(signs)
    assert x == x2


# ------------------------------------------------------------------ binding
def test_bind_xor_semantics():
    # bind(a,b) XOR semantics: bind(a, bind(a, b)) == b
    rng = np.random.default_rng(2)
    a = hv.random_hv(D=256, rng=rng)
    b = hv.random_hv(D=256, rng=rng)
    assert hv.bind(a, hv.bind(a, b)) == b


def test_unbind_equals_bind():
    rng = np.random.default_rng(3)
    a, b = hv.random_hv(D=64, rng=rng), hv.random_hv(D=64, rng=rng)
    assert hv.unbind(hv.bind(a, b), a) == hv.bind(hv.bind(a, b), a)


def test_bind_self_inverse_recovers_value():
    rng = np.random.default_rng(4)
    key = hv.random_hv(D=2000, rng=rng)
    value = hv.random_hv(D=2000, rng=rng)
    pair = hv.bind(key, value)
    recovered = hv.unbind(pair, key)
    # recovered HV should be *very* close to value (exact, since XOR is lossless)
    assert recovered == value


def test_bind_dimension_mismatch():
    a = hv.random_hv(D=100)
    b = hv.random_hv(D=200)
    with pytest.raises(ValueError):
        hv.bind(a, b)


# ----------------------------------------------------------------- bundling
def test_bundle_majority():
    # Three identical -> same; two vs one -> majority
    rng = np.random.default_rng(5)
    a = hv.random_hv(D=300, rng=rng)
    b = hv.random_hv(D=300, rng=rng)
    c = hv.random_hv(D=300, rng=rng)
    # a ⊕ a ⊕ b == a (a wins the majority)
    out = hv.bundle([a, a, b], rng=rng)
    assert out == a
    # a ⊕ a ⊕ a == a
    assert hv.bundle([a, a, a], rng=rng) == a


def test_bundle_with_weights():
    rng = np.random.default_rng(6)
    a = hv.random_hv(D=300, rng=rng)
    b = hv.random_hv(D=300, rng=rng)
    # weight b by 2 -> b wins
    out = hv.bundle([a, b], weights=[1.0, 2.0], rng=rng)
    assert out == b


def test_bundle_tie_random_break():
    rng = np.random.default_rng(7)
    a = hv.random_hv(D=400, rng=rng)
    # a ⊕ a_negated -> all ties -> output is random, but still a valid HV
    a_neg = hv.signs_to_bits(-hv.bits_to_signs(a))
    out = hv.bundle([a, a_neg], rng=rng)
    s = hv.bits_to_signs(out)
    assert set(np.unique(s).tolist()) <= {1, -1}


def test_bundle_empty_raises():
    with pytest.raises(ValueError):
        hv.bundle([])


def test_bundle_negative_weights_raise():
    a = hv.random_hv(D=10)
    b = hv.random_hv(D=10)
    with pytest.raises(ValueError):
        hv.bundle([a, b], weights=[-1.0, 1.0])


def test_bundle_dimension_mismatch():
    a = hv.random_hv(D=10)
    b = hv.random_hv(D=20)
    with pytest.raises(ValueError):
        hv.bundle([a, b])


# ----------------------------------------------------------- similarity
def test_similarity_self_is_one():
    x = hv.random_hv(D=1000)
    assert hv.similarity(x, x) == pytest.approx(1.0)


def test_similarity_opposite_is_minus_one():
    x = hv.random_hv(D=1000)
    neg = hv.signs_to_bits(-hv.bits_to_signs(x))
    assert hv.similarity(x, neg) == pytest.approx(-1.0)


def test_random_hvs_quasi_orthogonal():
    # Concentration of measure: with D large, two random HVs have
    # similarity concentrated around 0 with small std.
    rng = np.random.default_rng(8)
    sims = [hv.similarity(hv.random_hv(D=5000, rng=rng),
                          hv.random_hv(D=5000, rng=rng))
            for _ in range(200)]
    mean = float(np.mean(sims))
    std = float(np.std(sims))
    # mean near 0, std ~ 1/sqrt(D)
    assert abs(mean) < 0.02
    assert std < 1 / np.sqrt(5000) * 5  # generous bound


def test_hamming_basic():
    rng = np.random.default_rng(9)
    a = hv.random_hv(D=100, rng=rng)
    assert hv.hamming(a, a) == 0
    b = hv.random_hv(D=100, rng=rng)
    h = hv.hamming(a, b)
    assert 0 <= h <= 100
    assert h == hv.hamming(b, a)


def test_similarity_range():
    rng = np.random.default_rng(10)
    a = hv.random_hv(D=1000, rng=rng)
    b = hv.random_hv(D=1000, rng=rng)
    s = hv.similarity(a, b)
    assert -1.0 - 1e-9 <= s <= 1.0 + 1e-9


# ----------------------------------------------------------- hash / eq
def test_hash_equality_stable():
    rng = np.random.default_rng(11)
    a = hv.random_hv(D=100, rng=rng)
    a_copy = hv.HV(bits=a.bits.copy(), D=a.D)
    assert a == a_copy
    assert hash(a) == hash(a_copy)