tests/test_expert.py

"""Tests for the Expert: dataset -> compressed .exp -> reload round-trip."""
from __future__ import annotations

import os
import tempfile

import numpy as np
import pytest

from ensemble import Expert


# Small shared corpora for fast tests (D kept small for speed).
# Made large enough that fixed overhead doesn't dominate compression.
MATH_TEXT = (
    "what is two plus two. two plus two equals four. "
    "what is three times three. three times three equals nine. "
    "what is pi. pi is approximately three point one four. "
    "what is ten minus four. ten minus four equals six. "
    "what is the square root of nine. the square root of nine is three. "
) * 6

GEO_TEXT = (
    "the capital of france is paris. "
    "the capital of japan is tokyo. "
    "the capital of italy is rome. "
    "the capital of egypt is cairo. "
    "the capital of brazil is brasilia. "
) * 6

MATH_QA = [
    ("what is two plus two", "two plus two equals four"),
    ("what is three times three", "three times three equals nine"),
    ("what is pi", "pi is approximately three point one four"),
    ("what is ten minus four", "ten minus four equals six"),
    ("what is the square root of nine", "the square root of nine is three"),
    ("what is five times five", "five times five equals twenty five"),
    ("what is one hundred divided by ten", "one hundred divided by ten equals ten"),
    ("what is eight plus seven", "eight plus seven equals fifteen"),
] * 2


class TestExpertBuild:
    def test_from_text(self):
        e = Expert.from_text(MATH_TEXT, domain="math", D=2000)
        assert e.domain == "math"
        assert e.D == 2000
        assert e.n_traces > 0
        assert e.vocab_size > 4  # specials + chars

    def test_from_qa_pairs(self):
        e = Expert.from_qa_pairs(MATH_QA, domain="qa", D=2000)
        assert e.n_traces > 0
        assert e.domain == "qa"

    def test_from_text_records_source_stream(self):
        e = Expert.from_text(MATH_TEXT, domain="t", D=2000)
        assert e._source_token_stream is not None
        assert len(e._source_token_stream) > 0

    def test_qa_records_source_pairs(self):
        e = Expert.from_qa_pairs(MATH_QA, domain="q", D=2000)
        assert e._source_qa_pairs is not None
        assert len(e._source_qa_pairs) == len(MATH_QA)

    def test_signature_is_deterministic(self):
        e1 = Expert.from_text(MATH_TEXT, domain="m", D=2000, seed=0)
        e2 = Expert.from_text(MATH_TEXT, domain="m", D=2000, seed=0)
        assert e1.signature_hv == e2.signature_hv

    def test_different_seeds_give_different_self_hv(self):
        e1 = Expert.from_text(MATH_TEXT, domain="m", D=2000, seed=1)
        e2 = Expert.from_text(MATH_TEXT, domain="m", D=2000, seed=2)
        assert e1.model._self_hv != e2.model._self_hv


class TestExpertCompression:
    """The defining property: .exp is SMALLER than the source dataset."""

    def test_text_expert_smaller_than_source(self, tmp_path):
        e = Expert.from_text(MATH_TEXT, domain="math", D=2000)
        result = e.save(tmp_path / "math.exp")
        assert result.expert_size_bytes < result.source_size_bytes
        assert result.compression_ratio > 1.0

    def test_qa_expert_smaller_than_source(self, tmp_path):
        e = Expert.from_qa_pairs(MATH_QA, domain="qa", D=2000)
        result = e.save(tmp_path / "qa.exp")
        assert result.expert_size_bytes < result.source_size_bytes

    def test_expert_directory_structure(self, tmp_path):
        e = Expert.from_text(MATH_TEXT, domain="m", D=2000)
        path = tmp_path / "m.exp"
        result = e.save(path)
        p = type(path)(result.path)
        assert (p / "manifest.json").exists()
        assert (p / "vocab.json").exists()
        # lm mode has tokens.bin.gz, qa mode has qa_pairs.json.gz
        assert (p / "tokens.bin.gz").exists()


class TestExpertRoundTrip:
    """A saved expert reloads to a bit-identical model."""

    def test_text_round_trip(self, tmp_path):
        e = Expert.from_text(MATH_TEXT, domain="math", D=2000, seed=5)
        result = e.save(tmp_path / "math.exp")
        e2 = Expert.load(result.path)
        assert e2.D == e.D
        assert e2.domain == e.domain
        assert e2.n_traces == e.n_traces
        assert e2.signature_hv == e.signature_hv
        assert e2.model._self_hv == e.model._self_hv

    def test_qa_round_trip(self, tmp_path):
        e = Expert.from_qa_pairs(MATH_QA, domain="qa", D=2000, seed=3)
        result = e.save(tmp_path / "qa.exp")
        e2 = Expert.load(result.path)
        assert e2.n_traces == e.n_traces
        assert e2.signature_hv == e.signature_hv

    def test_manifest_preserved(self, tmp_path):
        e = Expert.from_text(MATH_TEXT, domain="science", D=2000)
        result = e.save(tmp_path / "s.exp")
        e2 = Expert.load(result.path)
        assert e2.manifest.domain == "science"
        assert e2.manifest.D == 2000
        assert e2.manifest.compression_ratio > 1.0


class TestExpertQuery:
    def test_answer_returns_string(self):
        e = Expert.from_qa_pairs(MATH_QA, domain="m", D=5000)
        a = e.answer("what is pi")
        assert isinstance(a, str)

    def test_relevance_returns_float(self):
        e = Expert.from_qa_pairs(MATH_QA, domain="m", D=2000)
        r = e.relevance("what is pi")
        assert -1.0 <= r <= 1.0

    def test_candidate_hv_returns_hv(self):
        from palimseste.hv import HV
        e = Expert.from_qa_pairs(MATH_QA, domain="m", D=2000)
        c = e.candidate_hv("what is pi")
        assert c is None or isinstance(c, HV)


class TestExpertDatasetFormats:
    def test_csv_qa(self, tmp_path):
        import csv as _csv
        p = tmp_path / "data.csv"
        with open(p, "w", newline="", encoding="utf-8") as f:
            w = _csv.writer(f)
            w.writerow(["question", "answer"])
            for q, a in MATH_QA:
                w.writerow([q, a])
        e = Expert.from_dataset(p, D=2000)
        assert e.n_traces > 0

    def test_json_qa(self, tmp_path):
        import json
        p = tmp_path / "data.json"
        with open(p, "w", encoding="utf-8") as f:
            json.dump([{"question": q, "answer": a} for q, a in MATH_QA], f)
        e = Expert.from_dataset(p, D=2000)
        assert e.n_traces > 0

    def test_jsonl(self, tmp_path):
        import json
        p = tmp_path / "data.jsonl"
        with open(p, "w", encoding="utf-8") as f:
            for q, a in MATH_QA:
                f.write(json.dumps({"question": q, "answer": a}) + "\n")
        e = Expert.from_dataset(p, D=2000)
        assert e.n_traces > 0

    def test_txt(self, tmp_path):
        p = tmp_path / "data.txt"
        p.write_text(MATH_TEXT, encoding="utf-8")
        e = Expert.from_dataset(p, D=2000)
        assert e.n_traces > 0