# tests/test_traces.py
from pathlib import Path

import numpy as np

from fractus_vorax.compiler.encode import text_hv
from fractus_vorax.hv import D_DEFAULT, sim
from fractus_vorax.organs.traces import TraceStore


def _fill(store: TraceStore) -> None:
    for i, q in enumerate(["what is the capital of france", "who wrote hamlet", "quantum stuff"]):
        addr, value = text_hv(q, D=2048), text_hv(q, D=2048)
        store.write(i, addr, value)


def test_write_and_retrieve_top1(tmp_path):
    store = TraceStore(D=2048)
    _fill(store)
    hits = store.retrieve(text_hv("what is the capital of france", D=2048), k=2)
    assert len(hits) == 2
    assert hits[0][0] == 0            # le bon atom_id en tête
    assert hits[0][1] > 0.5           # similarité élevée
    assert hits[0][1] >= hits[1][1]   # tri décroissant


def test_empty_retrieve_returns_empty():
    store = TraceStore(D=2048)
    assert store.retrieve(text_hv("rien", D=2048)) == []
    assert len(store) == 0


def test_weight_boosts_ranking():
    store = TraceStore(D=2048)
    q1, q2 = "what is the capital of france", "what is the capital of spain"
    store.write(0, text_hv(q1, D=2048), text_hv(q1, D=2048), weight=0.1)
    store.write(1, text_hv(q2, D=2048), text_hv(q2, D=2048), weight=2.0)
    hits = store.retrieve(text_hv("what is the capital of france", D=2048), k=2)
    assert hits[0][0] == 1  # poids fort domine malgré des tokens partagés


def test_save_load_roundtrip(tmp_path):
    store = TraceStore(D=2048)
    _fill(store)
    store.save(tmp_path / "brain")
    loaded = TraceStore.load(tmp_path / "brain")
    assert len(loaded) == 3
    q = text_hv("who wrote hamlet", D=2048)
    assert [i for i, _ in loaded.retrieve(q, k=3)] == [i for i, _ in store.retrieve(q, k=3)]