"""Tests for HV-Word2Vec, Soft Phi, and Multi-Level Query."""

from __future__ import annotations

import pytest
import numpy as np

from palimseste import hv
from palimseste.learner import Encoder
from palimseste.hv_word2vec import HVWord2Vec, Word2VecConfig
from palimseste.soft_phi import SoftPhi, SoftPhiConfig
from palimseste.multiquery import MultiLevelQuery, MultiLevelConfig
from palimseste.memory import Memory
from palimseste.phi import Phi, KernelConfig
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig


# ================================================================ HV-Word2Vec
class TestHVWord2Vec:
    def _train(self, text, D=3000, epochs=3):
        cfg = Word2VecConfig(D=D, n_epochs=epochs, min_word_freq=2)
        w2v = HVWord2Vec(config=cfg, rng=np.random.default_rng(42))
        w2v.train(text)
        return w2v

    def test_train_creates_embeddings(self):
        text = "the cat sat on the mat. the dog sat on the rug. the cat ran. the dog ran."
        w2v = self._train(text)
        assert w2v.vocab_size > 0
        assert w2v.is_trained

    def test_words_in_similar_contexts_are_similar(self):
        """cat and dog both appear with 'sat on the' → should be similar."""
        text = ("the cat sat on the mat. the dog sat on the rug. "
                "the cat sat on the bed. the dog sat on the floor. "
                "the cat ran fast. the dog ran fast. "
                "the cat ate food. the dog ate food.")
        w2v = self._train(text, D=5000, epochs=5)
        sim_cat_dog = w2v.similarity("cat", "dog")
        sim_cat_mat = w2v.similarity("cat", "mat")
        # cat and dog share contexts → more similar than cat and mat
        assert sim_cat_dog > sim_cat_mat, \
            f"cat-dog sim {sim_cat_dog:.3f} should be > cat-mat {sim_cat_mat:.3f}"

    def test_unknown_word_returns_none(self):
        w2v = self._train("hello world hello world")
        assert w2v.get_word_hv("xyzunknown") is None

    def test_encode_text_semantic(self):
        w2v = self._train("hello world hello world test test")
        h = w2v.encode_text_semantic("hello world")
        assert isinstance(h, hv.HV)
        assert h.D == 3000

    def test_most_similar(self):
        text = ("cat sat on mat. dog sat on mat. cat ran. dog ran. "
                "bird flew. bird sat on mat.")
        w2v = self._train(text, D=5000, epochs=5)
        similar = w2v.most_similar("cat", top_k=5)
        assert len(similar) > 0
        # dog should be in the top results (shared context with cat)
        words = [w for w, _ in similar]
        assert "dog" in words

    def test_save_load(self, tmp_path):
        w2v = self._train("test test hello hello world world")
        w2v.save(str(tmp_path / "w2v"))
        w2v2 = HVWord2Vec(config=Word2VecConfig(D=3000), rng=np.random.default_rng(0))
        w2v2.load(str(tmp_path / "w2v"))
        assert w2v2.vocab_size == w2v.vocab_size
        assert w2v2.get_word_hv("test") == w2v.get_word_hv("test")

    def test_config_validation(self):
        with pytest.raises(ValueError):
            Word2VecConfig(D=0)
        with pytest.raises(ValueError):
            Word2VecConfig(window_radius=0)


# ================================================================ Soft Phi
class TestSoftPhi:
    def _setup(self, D=3000):
        mem = Memory(D=D, rng=np.random.default_rng(0))
        rng = np.random.default_rng(1)
        # store 3 traces with known values
        for _ in range(3):
            a = hv.random_hv(D, rng=rng)
            v = hv.random_hv(D, rng=rng)
            mem.write(a, v)
        return mem, rng

    def test_soft_phi_returns_hv(self):
        mem, rng = self._setup()
        phi = SoftPhi(config=SoftPhiConfig(temperature=10.0))
        q = hv.random_hv(3000, rng=rng)
        result = phi(mem, q)
        # might be None if no candidates, or an HV
        if result is not None:
            assert isinstance(result, hv.HV)
            assert result.D == 3000

    def test_soft_phi_blends_results(self):
        """Soft Phi should produce a blend, not just the nearest match."""
        D = 5000
        mem = Memory(D=D, rng=np.random.default_rng(0))
        rng = np.random.default_rng(1)
        # create a query that's roughly between two stored addresses
        addr1 = hv.random_hv(D, rng=rng)
        addr2 = hv.random_hv(D, rng=rng)
        val1 = hv.random_hv(D, rng=rng)
        val2 = hv.random_hv(D, rng=rng)
        mem.write(addr1, val1)
        mem.write(addr2, val2)
        # query = midpoint (bundle of both addresses)
        query = hv.bundle([addr1, addr2], rng=rng, deterministic=True)

        phi = SoftPhi(config=SoftPhiConfig(temperature=5.0))
        result = phi(mem, query)
        if result is not None:
            # result should be somewhat similar to both val1 and val2
            sim1 = hv.similarity(result, val1)
            sim2 = hv.similarity(result, val2)
            # at least one should be positive (the blend favors the closer one)
            assert sim1 > 0 or sim2 > 0, \
                f"blend should resemble at least one value: sim1={sim1:.3f} sim2={sim2:.3f}"

    def test_soft_phi_empty_memory(self):
        mem = Memory(D=2000, rng=np.random.default_rng(0))
        phi = SoftPhi()
        q = hv.random_hv(2000, rng=np.random.default_rng(0))
        assert phi(mem, q) is None

    def test_soft_phi_retrieve(self):
        mem, rng = self._setup()
        phi = SoftPhi()
        q = hv.random_hv(3000, rng=rng)
        ret = phi.retrieve(mem, q)
        assert ret is not None  # Retrieval object
        assert isinstance(ret.matches, list)


# ================================================================ Multi-Level Query
class TestMultiLevelQuery:
    def _setup(self, D=3000):
        # train word2vec on a small corpus
        text = ("the capital of france is paris. "
                "the capital of italy is rome. "
                "the capital of germany is berlin. "
                "france italy germany are countries. "
                "paris rome berlin are cities. "
                "capital city country france italy germany paris rome berlin. "
                "what is the capital of france. "
                "what is the capital of italy. "
                "what is the capital of germany.")
        w2v = HVWord2Vec(config=Word2VecConfig(D=D, n_epochs=5, min_word_freq=2),
                         rng=np.random.default_rng(42))
        w2v.train(text)

        # create memory with traces
        enc = Encoder(D=D, rng=np.random.default_rng(42))
        mem = Memory(D=D, rng=np.random.default_rng(42))
        self_hv = hv.random_hv(D, rng=np.random.default_rng(42))

        # store fact: "capital of france" -> "paris"
        # at all three levels
        for question, answer in [("capital of france", "paris"),
                                  ("capital of italy", "rome"),
                                  ("capital of germany", "berlin")]:
            # surface level
            q_hvs = [enc.encode_str(w) for w in question.split()]
            q_bound = []
            for i, h in enumerate(q_hvs):
                q_bound.append(hv.bind(h, enc._role(i)))
            q_surface = hv.bundle(q_bound, rng=np.random.default_rng(0), deterministic=True)
            a_hv = enc.encode_str(f"__word_{answer}")
            mem.write(hv.bind(self_hv, q_surface), a_hv)

            # semantic level
            q_sem = w2v.encode_text_semantic(question)
            mem.write(hv.bind(self_hv, q_sem), a_hv)

        return mem, enc, w2v, self_hv

    def test_multi_level_query_returns_result(self):
        mem, enc, w2v, self_hv = self._setup()
        mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
                              self_hv=self_hv)
        # create a simple tokenizer-like object
        class FakeTok:
            def encode(self, text, add_bos=False, add_eos=False):
                return text.split()
            def token_hv(self, tid):
                return enc.encode_str(str(tid))
        
        phi = Phi(config=KernelConfig(radius=500, min_weight=1e-6))
        result = mlq.query("capital of france", FakeTok(), phi)
        assert result is not None
        assert result.source in ("surface", "semantic", "word", "none")

    def test_semantic_level_finds_related(self):
        """The semantic level should find 'capital of spain' related to
        stored 'capital of france' because the words are semantically similar."""
        mem, enc, w2v, self_hv = self._setup()
        mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
                              self_hv=self_hv)

        class FakeTok:
            def encode(self, text, add_bos=False, add_eos=False):
                return text.split()
            def token_hv(self, tid):
                return enc.encode_str(str(tid))

        phi = Phi(config=KernelConfig(radius=500, min_weight=1e-6))
        # "capital of spain" was never stored, but "capital" and "of" are
        # semantically related to stored "capital of france"
        result = mlq.query("capital of spain", FakeTok(), phi)
        # semantic level should have non-zero confidence
        assert result.semantic_confidence >= 0.0

    def test_three_encodings_different(self):
        """The three encodings should produce different HVs."""
        mem, enc, w2v, self_hv = self._setup()
        mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
                              self_hv=self_hv)

        class FakeTok:
            def encode(self, text, add_bos=False, add_eos=False):
                return text.split()
            def token_hv(self, tid):
                return enc.encode_str(str(tid))

        s = mlq.encode_surface("capital of france", FakeTok())
        w = mlq.encode_words("capital of france")
        sem = mlq.encode_semantic("capital of france")
        # all three should be valid HVs
        assert isinstance(s, hv.HV)
        assert isinstance(w, hv.HV)
        assert isinstance(sem, hv.HV)
        # they should not be identical
        assert s != w or w != sem