import json
import math
import os
import struct
import tempfile
import unittest

from hdc.hypervec import sim
from hdc.semantic import SemanticEncoder, StaticEmbeddings, available, load

D = 2000


def write_model(directory: str, vocab: dict[str, int], rows: list[list[float]]) -> str:
    """Write a minimal safetensors + tokenizer.json pair, no libraries."""
    flat = [value for row in rows for value in row]
    payload = struct.pack("<%df" % len(flat), *flat)
    header = json.dumps({
        "embeddings": {"dtype": "F32", "shape": [len(rows), len(rows[0])],
                       "data_offsets": [0, len(payload)]}
    }).encode("utf-8")
    header += b" " * ((8 - len(header) % 8) % 8)
    with open(os.path.join(directory, "model.safetensors"), "wb") as fh:
        fh.write(struct.pack("<Q", len(header)))
        fh.write(header)
        fh.write(payload)
    with open(os.path.join(directory, "tokenizer.json"), "w", encoding="utf-8") as fh:
        json.dump({"model": {"vocab": vocab, "type": "WordPiece",
                             "continuing_subword_prefix": "##"}}, fh)
    return directory


class SyntheticModel(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tmp = tempfile.TemporaryDirectory()
        vocab = {"[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "[MASK]": 4,
                 "cat": 5, "##s": 6, "car": 7, "auto": 8, "zzz": 9}
        rows = [
            [0.0] * 8, [0.1] * 8, [0.0] * 8, [0.0] * 8, [0.0] * 8,
            [1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],   # cat
            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],   # ##s
            [1.0, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],   # car
            [0.9, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],   # auto
            [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],   # zzz
        ]
        cls.path = write_model(cls.tmp.name, vocab, rows)
        cls.model = StaticEmbeddings(cls.path)

    @classmethod
    def tearDownClass(cls):
        cls.tmp.cleanup()

    def test_reads_the_matrix_shape(self):
        self.assertEqual(self.model.shape, [10, 8])
        self.assertEqual(self.model.dim, 8)
        self.assertEqual(len(self.model), 10)

    def test_greedy_wordpiece(self):
        self.assertEqual(self.model.tokenize("cats"), ["cat", "##s"])
        self.assertEqual(self.model.tokenize("car"), ["car"])

    def test_unknown_word_falls_back(self):
        self.assertIn("[UNK]", self.model.tokenize("qqqq"))

    def test_vectors_are_unit_length(self):
        vector = self.model.vector("cat")
        self.assertAlmostEqual(math.sqrt(sum(v * v for v in vector)), 1.0, places=6)

    def test_related_rows_are_more_similar(self):
        def dot(a, b):
            return sum(x * y for x, y in zip(a, b))

        self.assertGreater(dot(self.model.vector("car"), self.model.vector("auto")),
                           dot(self.model.vector("car"), self.model.vector("zzz")))

    def test_available_and_load(self):
        self.assertTrue(available(self.path))
        self.assertIsNotNone(load(D, self.path))
        self.assertFalse(available(os.path.join(self.path, "nope")))
        self.assertIsNone(load(D, os.path.join(self.path, "nope")))


class TestProjection(SyntheticModel):
    def test_encoding_is_deterministic(self):
        encoder = SemanticEncoder(path=self.path, d=D)
        self.assertEqual(encoder.encode("car"), encoder.encode("car"))

    def test_encoding_stays_inside_the_dimension(self):
        encoder = SemanticEncoder(path=self.path, d=D)
        self.assertLess(encoder.encode("cat").bit_length(), D + 1)

    def test_related_words_project_close(self):
        encoder = SemanticEncoder(path=self.path, d=D)
        related = sim(encoder.encode("car"), encoder.encode("auto"), D)
        unrelated = sim(encoder.encode("car"), encoder.encode("zzz"), D)
        self.assertGreater(related, unrelated)
        self.assertGreater(related, 0.2)

    def test_recovered_cosine_tracks_the_true_one(self):
        encoder = SemanticEncoder(path=self.path, d=D)
        true = encoder.similarity("car", "auto")
        recovered = encoder.cosine("car", "auto")
        self.assertLess(abs(recovered - true), 0.15)

    def test_projection_shrinks_high_cosines(self):
        # documented arcsine law: sim = (2/pi) arcsin(cos), so binarising
        # always reports less than the underlying cosine
        encoder = SemanticEncoder(path=self.path, d=D)
        raw = sim(encoder.encode("car"), encoder.encode("auto"), D)
        self.assertLess(raw, encoder.similarity("car", "auto"))


if __name__ == "__main__":
    unittest.main()