"""Tests for cortex features 4-7: multimodal, meta-learning, dual memory, analogies."""

import pytest
import numpy as np

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.cortex import (
    MultimodalFusion, MetaLearner, DualMemory, Analogizer,
    ModalityBinding, MetaLearningResult, MemoryRecord, AnalogyResult,
)
from palimseste.hv_word2vec import HVWord2Vec, Word2VecConfig


def _build_model(D=3000, ctx=64, radius=100):
    cfg = PalimpsesteConfig(D=D, context_window=ctx, kernel_radius=radius, temperature=0.0)
    lm = PalimpsesteForCausalLM(config=cfg)
    lm.build_tokenizer("hello world test memory image text analogical reasoning python france")
    return lm


# ================================================================ MULTIMODAL
class TestMultimodalFusion:
    def test_learn_image(self):
        lm = _build_model()
        fusion = MultimodalFusion(mem=lm.mem, encoder=lm.encoder)
        img = np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
        binding = fusion.learn_image("a red square", img)
        assert isinstance(binding, ModalityBinding)
        assert fusion.n_bindings == 1

    def test_find_image(self):
        lm = _build_model()
        fusion = MultimodalFusion(mem=lm.mem, encoder=lm.encoder)
        img = np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
        fusion.learn_image("a red square", img)
        results = fusion.find_image("red")
        assert len(results) > 0
        assert "red square" in results[0][0]

    def test_describe_image(self):
        lm = _build_model()
        fusion = MultimodalFusion(mem=lm.mem, encoder=lm.encoder)
        img = np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
        fusion.learn_image("a blue circle", img)
        results = fusion.describe_image(img)
        assert len(results) > 0
        # The exact image should match at sim ~1.0
        assert results[0][1] > 0.9

    def test_multiple_bindings(self):
        lm = _build_model()
        fusion = MultimodalFusion(mem=lm.mem, encoder=lm.encoder)
        for label in ["cat", "dog", "bird"]:
            img = np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
            fusion.learn_image(label, img)
        assert fusion.n_bindings == 3


# ================================================================ META-LEARNING
class TestMetaLearner:
    def test_adapt(self):
        lm = _build_model()
        # Train some data so there's a replay set
        lm.train_on_qa_pairs([("hello", "hi there"), ("what is x", "x is y")])
        meta = MetaLearner(mem=lm.mem, phi=lm.phi)
        result = meta.adapt()
        assert isinstance(result, MetaLearningResult)

    def test_current_config(self):
        lm = _build_model()
        meta = MetaLearner(mem=lm.mem, phi=lm.phi)
        cfg = meta.current_config
        assert "radius" in cfg
        assert "min_weight" in cfg

    def test_n_adaptations(self):
        lm = _build_model()
        meta = MetaLearner(mem=lm.mem, phi=lm.phi)
        assert meta.n_adaptations >= 0


# ================================================================ DUAL MEMORY
class TestDualMemory:
    def test_store_episodic(self):
        lm = _build_model()
        dm = DualMemory(mem=lm.mem, encoder=lm.encoder)
        record = dm.store_episodic("user said hello")
        assert record.memory_type == "episodic"
        assert dm.n_episodic == 1

    def test_store_semantic(self):
        lm = _build_model()
        dm = DualMemory(mem=lm.mem, encoder=lm.encoder)
        record = dm.store_semantic("the sky is blue")
        assert record.memory_type == "semantic"
        assert dm.n_semantic == 1

    def test_recall(self):
        lm = _build_model()
        dm = DualMemory(mem=lm.mem, encoder=lm.encoder)
        dm.store_semantic("python is a programming language")
        dm.store_episodic("user asked about python")
        results = dm.recall("python", top_k=3)
        assert len(results) > 0
        assert all(isinstance(r[0], MemoryRecord) for r in results)

    def test_forget_old(self):
        lm = _build_model()
        dm = DualMemory(mem=lm.mem, encoder=lm.encoder)
        dm.store_episodic("old memory")
        dm.store_semantic("persistent fact")
        # max_age=0 removes all episodic
        n = dm.forget_old_episodic(max_age=0)
        assert n >= 1
        assert dm.n_episodic == 0
        assert dm.n_semantic == 1

    def test_counts(self):
        lm = _build_model()
        dm = DualMemory(mem=lm.mem, encoder=lm.encoder)
        dm.store_episodic("e1")
        dm.store_episodic("e2")
        dm.store_semantic("s1")
        assert dm.n_episodic == 2
        assert dm.n_semantic == 1
        assert dm.total == 3


# ================================================================ ANALOGIES
class TestAnalogizer:
    def _build_w2v(self):
        w2v = HVWord2Vec(config=Word2VecConfig(D=3000, n_epochs=3, min_word_freq=1))
        corpus = ("france paris england london germany berlin italy rome japan tokyo "
                  "france paris england london germany berlin italy rome japan tokyo "
                  "france england germany italy japan paris london berlin rome tokyo")
        w2v.train(corpus)
        return w2v

    def test_analogy_returns_result(self):
        w2v = self._build_w2v()
        analogizer = Analogizer(word2vec=w2v)
        result = analogizer.analogy("paris", "france", "tokyo")
        assert result is not None
        assert isinstance(result, AnalogyResult)
        assert result.c == "tokyo"

    def test_analogy_unknown_word(self):
        w2v = self._build_w2v()
        analogizer = Analogizer(word2vec=w2v)
        result = analogizer.analogy("paris", "france", "xyznonexistent")
        assert result is None

    def test_analogy_no_w2v(self):
        analogizer = Analogizer(word2vec=None)
        result = analogizer.analogy("a", "b", "c")
        assert result is None

    def test_find_relations(self):
        w2v = self._build_w2v()
        analogizer = Analogizer(word2vec=w2v)
        related = analogizer.find_relations("france", top_k=3)
        assert len(related) > 0
        # Countries should be related to countries
        assert related[0][0] in ["italy", "england", "germany", "japan"]

    def test_find_relations_no_w2v(self):
        analogizer = Analogizer(word2vec=None)
        result = analogizer.find_relations("test")
        assert result == []