"""Tests for the Brain's central persistent memory."""
from __future__ import annotations

import pytest

from ensemble import Brain, Expert, BrainMemory
from palimseste.hv import random_hv, similarity


class TestBrainMemory:
    def test_create_empty(self):
        bm = BrainMemory(D=2000)
        assert bm.n_concepts == 0
        assert len(bm) == 0

    def test_remember_and_recall(self):
        bm = BrainMemory(D=2000)
        addr = random_hv(2000, rng=__import__("numpy").random.default_rng(1))
        val = random_hv(2000, rng=__import__("numpy").random.default_rng(2))
        bm.remember(addr, val, weight=1.0, tag="test")
        assert bm.n_concepts == 1
        # recall with the exact address -> should retrieve the value
        results = bm.recall(addr, top_k=1)
        assert len(results) >= 1
        assert results[0].value == val

    def test_recall_empty_memory(self):
        bm = BrainMemory(D=2000)
        q = random_hv(2000)
        assert bm.recall(q) == []

    def test_dimension_mismatch_raises(self):
        bm = BrainMemory(D=2000)
        bad = random_hv(3000)
        with pytest.raises(ValueError):
            bm.remember(bad, bad)

    def test_reconstruct(self):
        bm = BrainMemory(D=2000)
        rng = __import__("numpy").random.default_rng(5)
        addr = random_hv(2000, rng=rng)
        val = random_hv(2000, rng=rng)
        bm.remember(addr, val)
        recon = bm.reconstruct(addr)
        assert recon is not None
        assert similarity(recon, val) > 0.5

    def test_save_load_round_trip(self, tmp_path):
        bm = BrainMemory(D=2000)
        rng = __import__("numpy").random.default_rng(7)
        for i in range(5):
            bm.remember(random_hv(2000, rng=rng), random_hv(2000, rng=rng))
        assert bm.n_concepts == 5
        bm.save(tmp_path / "mem.bin")
        bm2 = BrainMemory.load(tmp_path / "mem.bin", D=2000)
        assert bm2.n_concepts == 5


class TestBrainCentralMemory:
    def test_brain_creates_memory_on_first_expert(self):
        brain = Brain()
        assert brain._brain_memory is None
        e = Expert.from_text("hello world. " * 20, D=2000)
        brain.add_expert(e)
        assert brain._brain_memory is not None
        assert brain._D == 2000

    def test_brain_remember_recall(self):
        brain = Brain()
        e = Expert.from_text("hello world. " * 20, D=2000)
        brain.add_expert(e)
        rng = __import__("numpy").random.default_rng(3)
        addr = random_hv(2000, rng=rng)
        val = random_hv(2000, rng=rng)
        brain.remember(addr, val, tag="manual")
        assert brain.n_concepts == 1
        results = brain.recall(addr)
        assert len(results) >= 1

    def test_brain_save_load_persists_concepts(self, tmp_path):
        brain = Brain()
        e = Expert.from_text("hello world. " * 20, D=2000)
        brain.add_expert(e)
        rng = __import__("numpy").random.default_rng(9)
        for _ in range(3):
            brain.remember(random_hv(2000, rng=rng), random_hv(2000, rng=rng))
        n_before = brain.n_concepts
        assert n_before == 3
        brain.save(tmp_path / "brain")
        brain2 = Brain.load(tmp_path / "brain")
        assert brain2.n_concepts == n_before

    def test_think_writes_to_central_memory(self):
        """thinking should write concepts into the central memory when they emerge."""
        # use overlapping experts so cross-expert concepts can form
        a = Expert.from_qa_pairs(
            [("what is x", "x is a thing")] * 4, domain="a", D=3000)
        b = Expert.from_qa_pairs(
            [("what is x", "x is an object")] * 4, domain="b", D=3000)
        brain = Brain()
        brain.add_expert(a)
        brain.add_expert(b)
        n0 = brain.n_concepts
        brain.think(n_cycles=4)
        # thinking ran without crashing; concepts may or may not form
        # (depends on coherence), but n_concepts must be accessible
        assert brain.n_concepts >= n0

    def test_brain_with_D_param(self):
        brain = Brain(D=2500)
        assert brain._D == 2500
        assert brain._brain_memory is not None
        assert brain._brain_memory.D == 2500