"""Tests for rbsp.core.quantize -- Binary and Product Quantization."""

from __future__ import annotations

import os
import random
import sys
import time
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from rbsp.core.quantize import (
    BinaryQuantizer,
    ProductQuantizer,
    hamming_distance,
    popcount,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class TestPopcount(unittest.TestCase):
    """popcount counts set bits in an integer."""

    def test_zero(self):
        self.assertEqual(popcount(0), 0)

    def test_one(self):
        self.assertEqual(popcount(1), 1)

    def test_all_ones_8bit(self):
        self.assertEqual(popcount(0xFF), 8)

    def test_power_of_two(self):
        self.assertEqual(popcount(1 << 63), 1)

    def test_alternating_bits(self):
        # 0b10101010 = 170
        self.assertEqual(popcount(0xAA), 4)

    def test_large_integer(self):
        # 128-bit number with all bits set
        val = (1 << 128) - 1
        self.assertEqual(popcount(val), 128)


class TestHammingDistance(unittest.TestCase):
    """hamming_distance via XOR + popcount."""

    def test_identical(self):
        self.assertEqual(hamming_distance(0xFF, 0xFF), 0)

    def test_one_bit_diff(self):
        self.assertEqual(hamming_distance(0b1110, 0b1111), 1)

    def test_all_bits_differ_8bit(self):
        self.assertEqual(hamming_distance(0x00, 0xFF), 8)

    def test_symmetric(self):
        a, b = 0b11001100, 0b10101010
        self.assertEqual(hamming_distance(a, b), hamming_distance(b, a))

    def test_zero_zero(self):
        self.assertEqual(hamming_distance(0, 0), 0)


# ---------------------------------------------------------------------------
# BinaryQuantizer
# ---------------------------------------------------------------------------


class TestBinaryQuantizer(unittest.TestCase):
    """BinaryQuantizer: float vectors -> binary codes -> Hamming search."""

    def test_quantize_without_fit(self):
        """Without fit, threshold is zero: positive -> 1."""
        bq = BinaryQuantizer()
        code = bq.quantize([1.0, -1.0, 0.5, -0.5])
        # bit 0: 1.0 >= 0 -> 1
        # bit 1: -1.0 < 0 -> 0
        # bit 2: 0.5 >= 0 -> 1
        # bit 3: -0.5 < 0 -> 0
        self.assertEqual(code, 0b0101)  # bits 0 and 2 set

    def test_dim_inferred(self):
        bq = BinaryQuantizer()
        self.assertIsNone(bq.dim)
        bq.quantize([1.0, 2.0, 3.0])
        self.assertEqual(bq.dim, 3)

    def test_dim_explicit(self):
        bq = BinaryQuantizer(dim=4)
        self.assertEqual(bq.dim, 4)

    def test_fit_computes_thresholds(self):
        bq = BinaryQuantizer()
        vectors = [
            [0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0],
            [2.0, 2.0, 2.0],
        ]
        bq.fit(vectors)
        self.assertIsNotNone(bq.thresholds)
        self.assertEqual(len(bq.thresholds), 3)
        # Median of [0, 1, 2] = 1.0
        for t in bq.thresholds:
            self.assertAlmostEqual(t, 1.0)

    def test_fit_empty_raises(self):
        bq = BinaryQuantizer()
        with self.assertRaises(ValueError):
            bq.fit([])

    def test_quantize_after_fit(self):
        bq = BinaryQuantizer()
        bq.fit([[0.0, 0.0], [1.0, 1.0]])
        # Threshold = 0.5 for both dims
        code = bq.quantize([0.6, 0.4])
        # bit 0: 0.6 >= 0.5 -> 1
        # bit 1: 0.4 < 0.5 -> 0
        self.assertEqual(code, 0b01)

    def test_insert_and_search(self):
        bq = BinaryQuantizer(dim=4)
        bq.insert("a", [1.0, 0.0, 1.0, 0.0])
        bq.insert("b", [0.0, 1.0, 0.0, 1.0])
        bq.insert("c", [1.0, 0.0, 1.0, 1.0])

        # Query closest to "a"
        results = bq.search([1.0, 0.0, 1.0, 0.0], k=2)
        self.assertEqual(len(results), 2)
        # "a" should be first (distance 0)
        self.assertEqual(results[0][0], "a")
        self.assertEqual(results[0][1], 0)

    def test_batch_insert(self):
        bq = BinaryQuantizer(dim=3)
        ids = ["x", "y", "z"]
        vecs = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
        bq.batch_insert(ids, vecs)
        self.assertEqual(bq.size, 3)

    def test_batch_insert_length_mismatch(self):
        bq = BinaryQuantizer()
        with self.assertRaises(ValueError):
            bq.batch_insert(["a", "b"], [[1.0]])

    def test_search_empty(self):
        bq = BinaryQuantizer(dim=2)
        results = bq.search([1.0, 0.0], k=5)
        self.assertEqual(results, [])

    def test_hamming_distance_method(self):
        """Static method matches module-level function."""
        self.assertEqual(BinaryQuantizer.hamming_distance(0b1010, 0b0110), 2)

    def test_size_property(self):
        bq = BinaryQuantizer(dim=2)
        self.assertEqual(bq.size, 0)
        bq.insert("a", [1.0, 0.0])
        self.assertEqual(bq.size, 1)


# ---------------------------------------------------------------------------
# ProductQuantizer
# ---------------------------------------------------------------------------


class TestProductQuantizer(unittest.TestCase):
    """ProductQuantizer: subspace k-means + ADC distance."""

    def _make_training_data(self, n=100, dim=8, seed=42):
        rng = random.Random(seed)
        return [[rng.gauss(0, 1) for _ in range(dim)] for _ in range(n)]

    def test_fit(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=50, dim=8)
        pq.fit(data, n_clusters=4)
        self.assertEqual(pq.dim, 8)
        self.assertEqual(pq.dsub, 4)
        self.assertEqual(pq.n_clusters, 4)

    def test_fit_empty_raises(self):
        pq = ProductQuantizer(m=2)
        with self.assertRaises(ValueError):
            pq.fit([])

    def test_fit_dim_not_divisible_raises(self):
        pq = ProductQuantizer(m=3)
        with self.assertRaises(ValueError):
            pq.fit([[1.0, 2.0, 3.0, 4.0, 5.0]])  # dim=5, m=3

    def test_encode_before_fit_raises(self):
        pq = ProductQuantizer(m=2)
        with self.assertRaises(RuntimeError):
            pq.encode([1.0, 2.0])

    def test_encode(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=50, dim=4)
        pq.fit(data, n_clusters=4)
        codes = pq.encode(data[0])
        self.assertEqual(len(codes), 2)
        for c in codes:
            self.assertGreaterEqual(c, 0)
            self.assertLess(c, 4)

    def test_decode(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=50, dim=4)
        pq.fit(data, n_clusters=4)
        codes = pq.encode(data[0])
        reconstructed = pq.decode(codes)
        self.assertEqual(len(reconstructed), 4)
        # Reconstruction should be approximate
        for r, o in zip(reconstructed, data[0]):
            self.assertIsInstance(r, float)

    def test_decode_before_fit_raises(self):
        pq = ProductQuantizer(m=2)
        with self.assertRaises(RuntimeError):
            pq.decode([0, 0])

    def test_distance_table(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=50, dim=4)
        pq.fit(data, n_clusters=4)
        table = pq.distance_table(data[0])
        self.assertEqual(len(table), 2)  # m subspaces
        for sub_table in table:
            self.assertEqual(len(sub_table), 4)  # n_clusters
            for d in sub_table:
                self.assertGreaterEqual(d, 0.0)

    def test_distance_table_before_fit_raises(self):
        pq = ProductQuantizer(m=2)
        with self.assertRaises(RuntimeError):
            pq.distance_table([1.0, 2.0])

    def test_insert_and_search(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=50, dim=4)
        pq.fit(data, n_clusters=4)

        for i, vec in enumerate(data):
            pq.insert(f"doc_{i}", vec)

        self.assertEqual(pq.size, 50)

        # Search for the first vector -- it should be top-1
        results = pq.search(data[0], k=5)
        self.assertEqual(len(results), 5)
        top_ids = [r[0] for r in results]
        self.assertIn("doc_0", top_ids)

    def test_search_returns_sorted(self):
        pq = ProductQuantizer(m=2)
        data = self._make_training_data(n=30, dim=4)
        pq.fit(data, n_clusters=4)
        for i, vec in enumerate(data):
            pq.insert(f"v_{i}", vec)

        results = pq.search(data[0], k=10)
        distances = [d for _, d in results]
        self.assertEqual(distances, sorted(distances))

    def test_properties_before_fit(self):
        pq = ProductQuantizer(m=4)
        self.assertEqual(pq.m, 4)
        self.assertIsNone(pq.dim)
        self.assertIsNone(pq.dsub)
        self.assertEqual(pq.size, 0)


# ---------------------------------------------------------------------------
# Performance
# ---------------------------------------------------------------------------


class TestPerformance(unittest.TestCase):
    """Hamming distance performance sanity check."""

    def test_100k_hamming_under_2s(self):
        """100K Hamming comparisons should complete in under 2 seconds."""
        rng = random.Random(42)
        codes_a = [rng.getrandbits(128) for _ in range(100_000)]
        codes_b = [rng.getrandbits(128) for _ in range(100_000)]

        t0 = time.perf_counter()
        for a, b in zip(codes_a, codes_b):
            hamming_distance(a, b)
        elapsed = time.perf_counter() - t0

        self.assertLess(elapsed, 2.0, f"100K Hamming took {elapsed:.2f}s")


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