"""Tests for rbsp.core.hnsw -- HNSW approximate nearest neighbor index."""

from __future__ import annotations

import math
import os
import random
import sys
import threading
import unittest

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

from rbsp.core.hnsw import (
    HNSWConfig,
    HNSWIndex,
    cosine_distance,
    dot_product,
    l2_distance,
    normalize_vector,
)


# ---------------------------------------------------------------------------
# Distance functions
# ---------------------------------------------------------------------------


class TestDotProduct(unittest.TestCase):
    """dot_product computes scalar inner product."""

    def test_basic(self):
        self.assertAlmostEqual(dot_product([1, 2, 3], [4, 5, 6]), 32.0)

    def test_orthogonal(self):
        self.assertAlmostEqual(dot_product([1, 0], [0, 1]), 0.0)

    def test_parallel(self):
        self.assertAlmostEqual(dot_product([3, 0], [3, 0]), 9.0)

    def test_zero_vector(self):
        self.assertAlmostEqual(dot_product([0, 0, 0], [1, 2, 3]), 0.0)

    def test_length_mismatch_raises(self):
        with self.assertRaises(ValueError):
            dot_product([1, 2], [1, 2, 3])


class TestNormalizeVector(unittest.TestCase):
    """normalize_vector returns unit-length copy."""

    def test_unit_length(self):
        v = normalize_vector([3, 4])
        norm = math.sqrt(sum(x * x for x in v))
        self.assertAlmostEqual(norm, 1.0)

    def test_direction_preserved(self):
        v = normalize_vector([3, 0])
        self.assertAlmostEqual(v[0], 1.0)
        self.assertAlmostEqual(v[1], 0.0)

    def test_zero_vector_unchanged(self):
        v = normalize_vector([0, 0, 0])
        self.assertEqual(v, [0, 0, 0])

    def test_already_unit(self):
        v = normalize_vector([1, 0, 0])
        self.assertAlmostEqual(v[0], 1.0)
        self.assertAlmostEqual(v[1], 0.0)
        self.assertAlmostEqual(v[2], 0.0)

    def test_negative_components(self):
        v = normalize_vector([-3, -4])
        norm = math.sqrt(sum(x * x for x in v))
        self.assertAlmostEqual(norm, 1.0)
        self.assertTrue(v[0] < 0)
        self.assertTrue(v[1] < 0)


class TestL2Distance(unittest.TestCase):
    """l2_distance computes Euclidean distance."""

    def test_same_point(self):
        self.assertAlmostEqual(l2_distance([1, 2, 3], [1, 2, 3]), 0.0)

    def test_basic(self):
        self.assertAlmostEqual(l2_distance([0, 0], [3, 4]), 5.0)

    def test_symmetry(self):
        a, b = [1, 2, 3], [4, 5, 6]
        self.assertAlmostEqual(l2_distance(a, b), l2_distance(b, a))

    def test_single_dimension(self):
        self.assertAlmostEqual(l2_distance([0], [7]), 7.0)

    def test_length_mismatch_raises(self):
        with self.assertRaises(ValueError):
            l2_distance([1, 2], [1])


class TestCosineDistance(unittest.TestCase):
    """cosine_distance computes 1 - cosine_similarity."""

    def test_identical_vectors(self):
        self.assertAlmostEqual(cosine_distance([1, 2, 3], [1, 2, 3]), 0.0)

    def test_scaled_vectors_same_direction(self):
        self.assertAlmostEqual(cosine_distance([1, 0], [100, 0]), 0.0)

    def test_opposite_vectors(self):
        self.assertAlmostEqual(cosine_distance([1, 0], [-1, 0]), 2.0)

    def test_orthogonal_vectors(self):
        self.assertAlmostEqual(cosine_distance([1, 0], [0, 1]), 1.0)

    def test_zero_vector_returns_one(self):
        self.assertAlmostEqual(cosine_distance([0, 0], [1, 2]), 1.0)
        self.assertAlmostEqual(cosine_distance([1, 2], [0, 0]), 1.0)

    def test_symmetry(self):
        a, b = [1, 2, 3], [4, 5, 6]
        self.assertAlmostEqual(cosine_distance(a, b), cosine_distance(b, a))

    def test_range(self):
        """Cosine distance should be in [0, 2]."""
        rng = random.Random(42)
        for _ in range(50):
            a = [rng.gauss(0, 1) for _ in range(8)]
            b = [rng.gauss(0, 1) for _ in range(8)]
            d = cosine_distance(a, b)
            self.assertGreaterEqual(d, -1e-9)
            self.assertLessEqual(d, 2.0 + 1e-9)

    def test_length_mismatch_raises(self):
        with self.assertRaises(ValueError):
            cosine_distance([1], [1, 2])


# ---------------------------------------------------------------------------
# HNSWConfig
# ---------------------------------------------------------------------------


class TestHNSWConfig(unittest.TestCase):
    """HNSWConfig defaults and validation."""

    def test_defaults(self):
        cfg = HNSWConfig()
        self.assertEqual(cfg.M, 16)
        self.assertEqual(cfg.ef_construction, 200)
        self.assertEqual(cfg.ef_search, 50)
        self.assertIsNone(cfg.max_level)
        self.assertEqual(cfg.distance, "cosine")
        self.assertEqual(cfg.seed, 42)

    def test_custom_values(self):
        cfg = HNSWConfig(
            M=32,
            ef_construction=400,
            ef_search=100,
            max_level=5,
            distance="l2",
            seed=99,
        )
        self.assertEqual(cfg.M, 32)
        self.assertEqual(cfg.ef_construction, 400)
        self.assertEqual(cfg.ef_search, 100)
        self.assertEqual(cfg.max_level, 5)
        self.assertEqual(cfg.distance, "l2")
        self.assertEqual(cfg.seed, 99)

    def test_invalid_M_raises(self):
        with self.assertRaises(ValueError):
            HNSWConfig(M=1)

    def test_invalid_ef_construction_raises(self):
        with self.assertRaises(ValueError):
            HNSWConfig(ef_construction=0)

    def test_invalid_ef_search_raises(self):
        with self.assertRaises(ValueError):
            HNSWConfig(ef_search=0)

    def test_invalid_distance_raises(self):
        with self.assertRaises(ValueError):
            HNSWConfig(distance="manhattan")

    def test_repr(self):
        cfg = HNSWConfig()
        r = repr(cfg)
        self.assertIn("HNSWConfig", r)
        self.assertIn("M=16", r)
        self.assertIn("cosine", r)


# ---------------------------------------------------------------------------
# HNSWIndex — basic operations
# ---------------------------------------------------------------------------


class TestHNSWIndexBasic(unittest.TestCase):
    """Insert, __len__, __contains__, __repr__."""

    def test_empty_index(self):
        idx = HNSWIndex()
        self.assertEqual(len(idx), 0)
        self.assertNotIn("x", idx)

    def test_insert_single(self):
        idx = HNSWIndex()
        idx.insert("a", [1.0, 0.0, 0.0])
        self.assertEqual(len(idx), 1)
        self.assertIn("a", idx)

    def test_insert_multiple(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        idx.insert("b", [0, 1])
        idx.insert("c", [1, 1])
        self.assertEqual(len(idx), 3)

    def test_insert_duplicate_replaces(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        idx.insert("a", [0, 1])
        self.assertEqual(len(idx), 1)
        self.assertIn("a", idx)
        # Should now be closest to [0, 1], not [1, 0]
        results = idx.search([0, 1], k=1)
        self.assertEqual(results[0][0], "a")
        self.assertAlmostEqual(results[0][1], 0.0, places=5)

    def test_repr(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        r = repr(idx)
        self.assertIn("HNSWIndex", r)
        self.assertIn("len=1", r)


# ---------------------------------------------------------------------------
# HNSWIndex — search
# ---------------------------------------------------------------------------


class TestHNSWIndexSearch(unittest.TestCase):
    """Search returns nearest neighbors."""

    def setUp(self):
        self.idx = HNSWIndex(HNSWConfig(distance="cosine", seed=42))
        self.idx.insert("right", [1, 0])
        self.idx.insert("up", [0, 1])
        self.idx.insert("diagonal", [1, 1])

    def test_exact_match(self):
        results = self.idx.search([1, 0], k=1)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0][0], "right")
        self.assertAlmostEqual(results[0][1], 0.0, places=5)

    def test_top_k(self):
        results = self.idx.search([1, 0], k=3)
        self.assertEqual(len(results), 3)
        # First result is the exact match
        self.assertEqual(results[0][0], "right")
        # Distances should be ascending
        for i in range(len(results) - 1):
            self.assertLessEqual(results[i][1], results[i + 1][1])

    def test_search_empty_index(self):
        idx = HNSWIndex()
        results = idx.search([1, 0], k=5)
        self.assertEqual(results, [])

    def test_search_k_larger_than_index(self):
        results = self.idx.search([1, 0], k=100)
        self.assertEqual(len(results), 3)

    def test_search_l2(self):
        idx = HNSWIndex(HNSWConfig(distance="l2", seed=42))
        idx.insert("origin", [0, 0])
        idx.insert("near", [0.1, 0.1])
        idx.insert("far", [10, 10])
        results = idx.search([0, 0], k=1)
        self.assertEqual(results[0][0], "origin")
        self.assertAlmostEqual(results[0][1], 0.0)


# ---------------------------------------------------------------------------
# HNSWIndex — remove
# ---------------------------------------------------------------------------


class TestHNSWIndexRemove(unittest.TestCase):
    """Remove nodes from the index."""

    def test_remove_existing(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        idx.insert("b", [0, 1])
        self.assertTrue(idx.remove("a"))
        self.assertEqual(len(idx), 1)
        self.assertNotIn("a", idx)
        self.assertIn("b", idx)

    def test_remove_nonexistent(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        self.assertFalse(idx.remove("z"))
        self.assertEqual(len(idx), 1)

    def test_remove_all(self):
        idx = HNSWIndex()
        idx.insert("a", [1, 0])
        idx.insert("b", [0, 1])
        idx.remove("a")
        idx.remove("b")
        self.assertEqual(len(idx), 0)
        self.assertEqual(idx.search([1, 0], k=1), [])

    def test_remove_entry_point(self):
        """Removing the entry point should not break the index."""
        idx = HNSWIndex(HNSWConfig(seed=42))
        idx.insert("a", [1, 0, 0])
        idx.insert("b", [0, 1, 0])
        idx.insert("c", [0, 0, 1])
        # Remove whatever the entry point is -- the index should still work
        idx.remove("a")
        results = idx.search([0, 1, 0], k=1)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0][0], "b")

    def test_search_after_remove(self):
        idx = HNSWIndex(HNSWConfig(distance="l2", seed=42))
        idx.insert("a", [0, 0])
        idx.insert("b", [1, 0])
        idx.insert("c", [10, 0])
        idx.remove("b")
        results = idx.search([0.5, 0], k=1)
        # b is gone; closest remaining is a (dist 0.5) or c (dist 9.5)
        self.assertEqual(results[0][0], "a")


# ---------------------------------------------------------------------------
# HNSWIndex — edge cases
# ---------------------------------------------------------------------------


class TestHNSWIndexEdgeCases(unittest.TestCase):
    """Edge cases: single element, high dimensions, etc."""

    def test_single_element_search(self):
        idx = HNSWIndex()
        idx.insert("only", [1, 2, 3])
        results = idx.search([1, 2, 3], k=5)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0][0], "only")

    def test_high_dimensional(self):
        """128-dimensional vectors should work."""
        rng = random.Random(42)
        idx = HNSWIndex(HNSWConfig(distance="l2", seed=42))
        for i in range(20):
            vec = [rng.gauss(0, 1) for _ in range(128)]
            idx.insert(f"v{i}", vec)
        self.assertEqual(len(idx), 20)
        query = [rng.gauss(0, 1) for _ in range(128)]
        results = idx.search(query, k=5)
        self.assertEqual(len(results), 5)

    def test_identical_vectors(self):
        idx = HNSWIndex(HNSWConfig(distance="l2", seed=42))
        for i in range(5):
            idx.insert(f"v{i}", [1, 1, 1])
        results = idx.search([1, 1, 1], k=5)
        self.assertEqual(len(results), 5)
        for _, dist in results:
            self.assertAlmostEqual(dist, 0.0)


# ---------------------------------------------------------------------------
# Thread safety
# ---------------------------------------------------------------------------


class TestHNSWThreadSafety(unittest.TestCase):
    """Concurrent inserts should not corrupt the index."""

    def test_concurrent_inserts(self):
        idx = HNSWIndex(HNSWConfig(seed=42))
        errors: list = []

        def inserter(start: int, count: int):
            rng = random.Random(start)
            try:
                for i in range(count):
                    vec = [rng.gauss(0, 1) for _ in range(8)]
                    idx.insert(f"t{start}_v{i}", vec)
            except Exception as exc:
                errors.append(exc)

        threads = []
        num_threads = 4
        per_thread = 25
        for t in range(num_threads):
            th = threading.Thread(target=inserter, args=(t, per_thread))
            threads.append(th)
            th.start()

        for th in threads:
            th.join()

        self.assertEqual(errors, [], f"Errors during concurrent insert: {errors}")
        self.assertEqual(len(idx), num_threads * per_thread)

    def test_concurrent_search_and_insert(self):
        """Mixing inserts and searches should not raise."""
        idx = HNSWIndex(HNSWConfig(seed=42))
        # Pre-populate
        rng = random.Random(42)
        for i in range(20):
            idx.insert(f"base_{i}", [rng.gauss(0, 1) for _ in range(8)])

        errors: list = []

        def searcher():
            rng_s = random.Random(99)
            try:
                for _ in range(20):
                    q = [rng_s.gauss(0, 1) for _ in range(8)]
                    idx.search(q, k=3)
            except Exception as exc:
                errors.append(exc)

        def inserter():
            rng_i = random.Random(77)
            try:
                for i in range(20):
                    idx.insert(f"new_{i}", [rng_i.gauss(0, 1) for _ in range(8)])
            except Exception as exc:
                errors.append(exc)

        t1 = threading.Thread(target=searcher)
        t2 = threading.Thread(target=inserter)
        t1.start()
        t2.start()
        t1.join()
        t2.join()

        self.assertEqual(errors, [], f"Errors during concurrent ops: {errors}")


# ---------------------------------------------------------------------------
# Search quality / recall
# ---------------------------------------------------------------------------


class TestHNSWRecall(unittest.TestCase):
    """Search quality: verify recall vs brute-force on random data."""

    def _brute_force_knn(self, vectors, query, k, dist_fn):
        """Exact k-NN via linear scan."""
        scored = []
        for vid, vec in vectors.items():
            d = dist_fn(query, vec)
            scored.append((d, vid))
        scored.sort(key=lambda x: x[0])
        return [vid for _, vid in scored[:k]]

    def test_recall_cosine_100_vectors(self):
        """100 random 32-dim vectors, 10 queries, recall > 80%."""
        dim = 32
        n = 100
        k = 10
        num_queries = 10
        rng = random.Random(42)

        cfg = HNSWConfig(
            M=16,
            ef_construction=200,
            ef_search=50,
            distance="cosine",
            seed=42,
        )
        idx = HNSWIndex(cfg)

        vectors = {}
        for i in range(n):
            vec = tuple(rng.gauss(0, 1) for _ in range(dim))
            vid = f"v{i}"
            vectors[vid] = vec
            idx.insert(vid, vec)

        total_recall = 0.0
        for q_idx in range(num_queries):
            query = tuple(rng.gauss(0, 1) for _ in range(dim))
            true_nn = set(self._brute_force_knn(vectors, query, k, cosine_distance))
            hnsw_results = idx.search(query, k=k)
            hnsw_nn = {nid for nid, _ in hnsw_results}
            recall = len(true_nn & hnsw_nn) / k
            total_recall += recall

        avg_recall = total_recall / num_queries
        self.assertGreater(
            avg_recall,
            0.80,
            f"Average recall {avg_recall:.2f} is below 80% threshold",
        )

    def test_recall_l2_100_vectors(self):
        """100 random 32-dim vectors with L2 distance, recall > 80%."""
        dim = 32
        n = 100
        k = 10
        num_queries = 10
        rng = random.Random(123)

        cfg = HNSWConfig(
            M=16,
            ef_construction=200,
            ef_search=50,
            distance="l2",
            seed=123,
        )
        idx = HNSWIndex(cfg)

        vectors = {}
        for i in range(n):
            vec = tuple(rng.gauss(0, 1) for _ in range(dim))
            vid = f"v{i}"
            vectors[vid] = vec
            idx.insert(vid, vec)

        total_recall = 0.0
        for q_idx in range(num_queries):
            query = tuple(rng.gauss(0, 1) for _ in range(dim))
            true_nn = set(self._brute_force_knn(vectors, query, k, l2_distance))
            hnsw_results = idx.search(query, k=k)
            hnsw_nn = {nid for nid, _ in hnsw_results}
            recall = len(true_nn & hnsw_nn) / k
            total_recall += recall

        avg_recall = total_recall / num_queries
        self.assertGreater(
            avg_recall,
            0.80,
            f"Average recall {avg_recall:.2f} is below 80% threshold",
        )

    def test_exact_nearest_neighbor(self):
        """The single nearest neighbor should almost always be correct."""
        dim = 16
        n = 50
        rng = random.Random(7)

        cfg = HNSWConfig(distance="l2", ef_search=100, seed=7)
        idx = HNSWIndex(cfg)

        vectors = {}
        for i in range(n):
            vec = tuple(rng.gauss(0, 1) for _ in range(dim))
            vid = f"v{i}"
            vectors[vid] = vec
            idx.insert(vid, vec)

        correct = 0
        num_queries = 20
        for _ in range(num_queries):
            query = tuple(rng.gauss(0, 1) for _ in range(dim))
            true_nn = self._brute_force_knn(vectors, query, 1, l2_distance)[0]
            hnsw_nn = idx.search(query, k=1)[0][0]
            if hnsw_nn == true_nn:
                correct += 1

        # At least 80% of 1-NN queries should be exact
        self.assertGreaterEqual(
            correct / num_queries,
            0.80,
            f"1-NN accuracy {correct}/{num_queries} is too low",
        )


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