"""Tests for the causal LM, tokenizer, serialization, and HF interface.

Verifies:
  - tokenizer encode/decode round-trips, vocab building
  - LM training is O(1) and grows M
  - LM next-token prediction achieves 100% accuracy on training data
    (exact-context retrieval with deterministic bundling)
  - generate() produces text and respects max_new_tokens / EOS
  - save_pretrained / from_pretrained round-trips bit-identically
  - config presets (small, 1b) have correct D values
  - model card generation
"""

from __future__ import annotations

import tempfile
import numpy as np
import pytest

from palimseste import hv
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig, PRESETS
from palimseste.tokenizer import CharTokenizer, VOCAB_SPECIAL, PAD, BOS, EOS, UNK
from palimseste.learner import Encoder
from palimseste.serialization import save_memory, load_memory, save_encoder, load_encoder
from palimseste.hf import HFPalimpsesteLM, generate_model_card


# ----------------------------------------------------------------- tokenizer
class TestTokenizer:
    def _enc(self, D=1000):
        return Encoder(D=D, rng=np.random.default_rng(0))

    def test_encode_decode_roundtrip(self):
        tok = CharTokenizer(encoder=self._enc())
        text = "hello world 123"
        tok.build_vocab(text)
        ids = tok.encode(text)
        assert tok.decode(ids) == text

    def test_special_tokens_present(self):
        tok = CharTokenizer(encoder=self._enc())
        assert tok.id2char[PAD] == "<pad>"
        assert tok.id2char[BOS] == "<bos>"
        assert tok.id2char[EOS] == "<eos>"
        assert tok.id2char[UNK] == "<unk>"

    def test_encode_with_bos_eos(self):
        tok = CharTokenizer(encoder=self._enc())
        tok.build_vocab("ab")
        ids = tok.encode("ab", add_bos=True, add_eos=True)
        assert ids[0] == BOS
        assert ids[-1] == EOS
        assert tok.decode(ids) == "ab"  # specials stripped on decode

    def test_unknown_char_maps_to_unk(self):
        tok = CharTokenizer(encoder=self._enc())
        tok.build_vocab("abc")
        ids = tok.encode("abcé")
        assert ids[-1] == UNK

    def test_vocab_size_grows(self):
        tok = CharTokenizer(encoder=self._enc())
        n0 = tok.vocab_size
        tok.build_vocab("abc")
        assert tok.vocab_size == n0 + 3

    def test_token_hv_stable(self):
        tok = CharTokenizer(encoder=self._enc())
        tok.build_vocab("ab")
        assert tok.token_hv(tok.char2id["a"]) == tok.token_hv(tok.char2id["a"])
        assert tok.token_hv(tok.char2id["a"]) != tok.token_hv(tok.char2id["b"])

    def test_save_load_vocabulary(self, tmp_path):
        enc = self._enc()
        tok = CharTokenizer(encoder=enc)
        tok.build_vocab("hello world")
        tok.save_vocabulary(tmp_path / "vocab.json")
        tok2 = CharTokenizer.load_vocabulary(tmp_path / "vocab.json", encoder=enc)
        assert tok2.vocab_size == tok.vocab_size
        assert tok2.id2char == tok.id2char
        assert tok2.encode("hello") == tok.encode("hello")


# ----------------------------------------------------------------- serialization
class TestSerialization:
    def test_memory_round_trip(self, tmp_path):
        from palimseste.memory import Memory
        mem = Memory(D=2000, rng=np.random.default_rng(0))
        rng = np.random.default_rng(1)
        for _ in range(50):
            mem.write(hv.random_hv(D=2000, rng=rng), hv.random_hv(D=2000, rng=rng))
        n_before = len(mem)
        save_memory(mem, tmp_path / "mem.bin")
        mem2 = load_memory(tmp_path / "mem.bin", rng=np.random.default_rng(99))
        assert len(mem2) == n_before
        # every trace's address and value must match
        for a, b in zip(mem.traces, mem2.traces):
            assert a.address == b.address
            assert a.value == b.value
            assert a.weight == b.weight

    def test_memory_candidates_preserved(self, tmp_path):
        """LSH index must produce the same candidates after reload."""
        from palimseste.memory import Memory
        mem = Memory(D=2000, rng=np.random.default_rng(0))
        rng = np.random.default_rng(1)
        addrs = [hv.random_hv(D=2000, rng=rng) for _ in range(30)]
        for a in addrs:
            mem.write(a, hv.random_hv(D=2000, rng=rng))
        save_memory(mem, tmp_path / "mem.bin")
        mem2 = load_memory(tmp_path / "mem.bin", rng=np.random.default_rng(99))
        for a in addrs[:5]:
            assert mem.candidates(a) == mem2.candidates(a)

    def test_encoder_round_trip(self, tmp_path):
        enc = Encoder(D=2000, rng=np.random.default_rng(0))
        # materialize some atoms, roles, and levels
        enc.encode_int(42)
        enc.encode_str("cat")
        enc.encode_float(0.5)
        enc.encode_sequence([enc.encode_int(1), enc.encode_int(2)])
        save_encoder(enc, tmp_path / "enc.json")
        enc2 = load_encoder(tmp_path / "enc.json", rng=np.random.default_rng(99))
        assert enc._atoms == enc2._atoms
        assert enc._roles == enc2._roles
        # encode_float uses levels; check they round-trip
        if enc._levels and enc2._levels:
            for a, b in zip(enc._levels, enc2._levels):
                assert a == b

    def test_encoder_int_role_keys_round_trip(self, tmp_path):
        """Role keys are plain ints; must survive JSON serialization."""
        enc = Encoder(D=500, rng=np.random.default_rng(0))
        enc.encode_sequence([enc.encode_str("a"), enc.encode_str("b"), enc.encode_str("c")])
        save_encoder(enc, tmp_path / "enc.json")
        enc2 = load_encoder(tmp_path / "enc.json")
        assert enc._roles == enc2._roles
        # all role keys must be ints, not strings
        for k in enc2._roles:
            assert isinstance(k, int), f"role key {k} is {type(k)}, expected int"


# ----------------------------------------------------------------- LM
class TestLM:
    def _lm(self, D=3000, radius=50):
        cfg = PalimpsesteConfig(D=D, context_window=10, kernel_radius=radius, temperature=0.3)
        return PalimpsesteForCausalLM(config=cfg)

    def test_train_grows_memory(self):
        lm = self._lm()
        text = "hello world"
        lm.build_tokenizer(text)
        n = lm.train_on_text(text)
        assert n > 0
        assert len(lm.mem) == n

    def test_next_token_accuracy_on_training_data(self):
        lm = self._lm(radius=0)
        text = "the quick brown fox jumps over the lazy dog"
        lm.build_tokenizer(text)
        lm.train_on_text(text)
        ev = lm.evaluate(text)
        # with radius=0 and deterministic bundling, exact contexts are recovered
        assert ev["next_token_accuracy"] > 0.9

    def test_generate_returns_text(self):
        lm = self._lm()
        text = "the quick brown fox jumps over the lazy dog. the lazy dog sleeps."
        lm.build_tokenizer(text)
        lm.train_on_text(text)
        out = lm.generate("the ", max_new_tokens=20, temperature=0.0, seed=0)
        assert isinstance(out.text, str)
        assert len(out.token_ids) <= 20

    def test_generate_respects_max_tokens(self):
        lm = self._lm()
        text = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
        lm.build_tokenizer(text)
        lm.train_on_text(text)
        out = lm.generate("a", max_new_tokens=5, temperature=0.5, seed=0)
        assert len(out.token_ids) <= 5

    def test_generate_cold_start_no_crash(self):
        # untrained model should not crash on generate
        lm = self._lm()
        lm.build_tokenizer("abc")
        out = lm.generate("a", max_new_tokens=10, temperature=0.5, seed=0)
        assert isinstance(out.text, str)

    def test_predict_next(self):
        lm = self._lm(radius=0)
        text = "the quick brown fox"
        lm.build_tokenizer(text)
        lm.train_on_text(text)
        ids = lm.tokenizer.encode(text, add_bos=True)
        ctx = ids[:3]  # [BOS, 't', 'h']
        tid, conf = lm.predict_next(ctx)
        assert 0 <= tid < lm.tokenizer.vocab_size
        assert 0.0 <= conf <= 1.0

    def test_stats(self):
        lm = self._lm()
        lm.build_tokenizer("abc")
        lm.train_on_text("abc")
        s = lm.stats()
        assert s["D"] == 3000
        assert s["n_traces"] > 0
        assert s["vocab_size"] > 0
        assert s["theoretical_capacity_log2"] > 0

    def test_presets(self):
        assert PRESETS["small"].D == 10_000
        assert PRESETS["1b"].D == 100_000
        assert PRESETS["tiny"].D == 2_000

    def test_config_roundtrip(self):
        cfg = PalimpsesteConfig(D=5000, context_window=20, kernel_radius=100, temperature=0.7)
        d = cfg.to_dict()
        cfg2 = PalimpsesteConfig.from_dict(d)
        assert cfg2 == cfg


# ----------------------------------------------------------------- HF interface
class TestHFInterface:
    def test_save_load_round_trip_identical_generation(self, tmp_path):
        cfg = PalimpsesteConfig(D=3000, context_window=10, kernel_radius=50, temperature=0.3)
        lm = HFPalimpsesteLM(config=cfg, rng=np.random.default_rng(42))
        text = "the quick brown fox jumps over the lazy dog. the lazy dog sleeps."
        lm.build_tokenizer(text)
        lm.train_on_text(text)
        out1 = lm.generate("the quick ", max_new_tokens=15, temperature=0.0, seed=0)

        lm.save_pretrained(tmp_path / "model")
        lm2 = HFPalimpsesteLM.from_pretrained(tmp_path / "model")
        out2 = lm2.generate("the quick ", max_new_tokens=15, temperature=0.0, seed=0)
        assert out1.text == out2.text

    def test_save_creates_expected_files(self, tmp_path):
        cfg = PalimpsesteConfig(D=2000, context_window=8, kernel_radius=50)
        lm = HFPalimpsesteLM(config=cfg)
        lm.build_tokenizer("hello")
        lm.train_on_text("hello")
        d = tmp_path / "model"
        lm.save_pretrained(d)
        files = set(p.name for p in d.iterdir())
        assert "config.json" in files
        assert "palimpseste_memory.bin" in files
        assert "vocab.json" in files
        assert "README.md" in files

    def test_model_card(self):
        cfg = PalimpsesteConfig(D=10000, context_window=32, kernel_radius=200, vocab_size=100)
        stats = {"theoretical_capacity_log2": 2500.0}
        card = generate_model_card(cfg, stats)
        assert "PALIMPSESTE" in card
        assert "10,000" in card
        assert "no weight matrix" in card.lower() or "no gradient" in card.lower()

    def test_config_json_loads(self, tmp_path):
        cfg = PalimpsesteConfig(D=4000, context_window=16, kernel_radius=120, temperature=0.5)
        lm = HFPalimpsesteLM(config=cfg)
        lm.build_tokenizer("test")
        lm.train_on_text("test")
        d = tmp_path / "model"
        lm.save_pretrained(d)
        lm2 = HFPalimpsesteLM.from_pretrained(d)
        assert lm2.config.D == 4000
        assert lm2.config.context_window == 16
        assert lm2.config.kernel_radius == 120