"""Tests for the structural query encoder (patterns + slots)."""
from __future__ import annotations

import pytest

from ensemble import Expert, Brain
from ensemble.structural import StructuralEncoder
from palimseste.learner import Encoder
from palimseste.hv import similarity
import numpy as np


@pytest.fixture
def encoder():
    return Encoder(D=3000, rng=np.random.default_rng(0))


@pytest.fixture
def se(encoder):
    s = StructuralEncoder(encoder, D=3000)
    s.learn_pattern("what is the capital of {country}")
    s.learn_pattern("what is the population of {country}")
    return s


class TestStructuralEncoder:
    def test_learn_pattern(self, se):
        assert se.n_patterns == 2

    def test_learn_invalid_template(self, encoder):
        se = StructuralEncoder(encoder, D=3000)
        with pytest.raises(ValueError):
            se.learn_pattern("no placeholder here")
        with pytest.raises(ValueError):
            se.learn_pattern("two {a} placeholders {b}")

    def test_learn_idempotent(self, se):
        n_before = se.n_patterns
        se.learn_pattern("what is the capital of {country}")
        assert se.n_patterns == n_before  # no duplicate

    def test_encode_matches(self, se):
        r = se.encode("what is the capital of france")
        assert r.matched
        assert r.pattern is not None
        assert r.slot_value == "france"
        assert r.slot_hv is not None

    def test_encode_no_match(self, se):
        r = se.encode("what color is the sky")
        assert not r.matched
        assert r.pattern is None

    def test_slot_hv_stable(self, se):
        r1 = se.encode("what is the capital of france")
        r2 = se.encode("what is the capital of france")
        assert r1.slot_hv == r2.slot_hv

    def test_observe_and_nearest_slot(self, se):
        p = se.patterns[0]
        se.observe_slot(p, "france")
        se.observe_slot(p, "germany")
        name, sim = se.nearest_slot(p, se._slot_hv("france"))
        assert name == "france"
        assert sim > 0.99  # identical

    def test_nearest_slot_empty(self, se):
        p = se.patterns[0]
        name, sim = se.nearest_slot(p, se._slot_hv("x"))
        assert name is None

    def test_to_from_dict_roundtrip(self, se, encoder):
        d = se.to_dict()
        se2 = StructuralEncoder.from_dict(d, encoder)
        assert se2.n_patterns == se.n_patterns
        assert [p.template for p in se2.patterns] == [p.template for p in se.patterns]

    def test_bind_cancels_pattern(self, se):
        """bind(pattern, slot) means same-pattern states differ only by slot."""
        r_fr = se.encode("what is the capital of france")
        r_sp = se.encode("what is the capital of spain")
        r_pop = se.encode("what is the population of france")
        # different patterns: orthogonal (bind with different pattern_hv)
        assert similarity(r_fr.state_hv, r_pop.state_hv) < 0.1


class TestExpertStructural:
    QA = [
        ("what is the capital of france", "the capital of france is paris"),
        ("what is the capital of germany", "the capital of germany is berlin"),
        ("what is the capital of italy", "the capital of italy is rome"),
    ] * 3
    PATTERNS = ["what is the capital of {country}"]

    def test_expert_learns_patterns(self):
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000,
                                 patterns=self.PATTERNS)
        assert e.n_patterns == 1
        assert e._structural is not None

    def test_relevance_positive_for_pattern(self):
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000,
                                 patterns=self.PATTERNS)
        # known AND unknown slots both match the pattern -> positive relevance
        assert e.relevance("what is the capital of france") > 0.3
        assert e.relevance("what is the capital of spain") > 0.3

    def test_unknown_slot_returns_nonempty(self):
        """THE generalization guarantee: unseen slot -> non-empty answer."""
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000,
                                 patterns=self.PATTERNS)
        # spain was never in training data
        ans = e.answer("what is the capital of spain")
        assert ans != "", "structural expert should return a non-empty answer for unseen slots"

    def test_char_level_returns_empty_for_unknown(self):
        """Baseline: char-level (no patterns) returns empty for unseen slots."""
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000)  # no patterns
        ans = e.answer("what is the capital of spain")
        # char-level generalizes poorly -> empty or near-empty
        assert ans == "" or len(ans) < 5

    def test_no_patterns_falls_back_to_char(self):
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000)
        assert e._structural is None
        # relevance uses char-level path, doesn't crash
        r = e.relevance("what is the capital of france")
        assert -1.0 <= r <= 1.0


class TestStructuralBrain:
    QA = [
        ("what is the capital of france", "the capital of france is paris"),
        ("what is the capital of germany", "the capital of germany is berlin"),
        ("what is the capital of italy", "the capital of italy is rome"),
    ] * 3

    def test_brain_routes_structural_expert(self):
        e = Expert.from_qa_pairs(self.QA, domain="geo", D=3000,
                                 patterns=["what is the capital of {country}"])
        brain = Brain()
        brain.add_expert(e)
        res = brain.query("what is the capital of spain")
        # the structural expert should be routed to (dominant) and answer non-empty
        assert res.dominant_expert == "geo"
        assert res.answer != ""