"""
Tests for Spatial Memory Tool

Verifies the FibPi3D-backed spatial memory system provides:
1. Deterministic scene generation (same seed = same scene)
2. Correct spatial query results vs brute-force ground truth
3. Receipt generation for all operations
"""

import hashlib
import math
import pytest
import numpy as np
from pathlib import Path

import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from helix_cdc.echo.ops.spatial_memory import (
    SpatialMemoryManager,
    SceneConfig,
    SceneObject,
    DeterministicRNG,
    generate_scene,
    compute_scene_hash,
    FibPi3DBackend,
)


# ============================================================================
# Tier 1: Determinism Tests
# ============================================================================

class TestDeterminism:
    """Verify deterministic behavior."""

    def test_rng_reproducibility(self):
        """Same seed produces same random sequence."""
        rng1 = DeterministicRNG("test_seed")
        rng2 = DeterministicRNG("test_seed")

        values1 = [rng1.random() for _ in range(100)]
        values2 = [rng2.random() for _ in range(100)]

        assert values1 == values2, "RNG should produce identical sequences for same seed"

    def test_scene_generation_determinism(self):
        """Same config produces same scene."""
        config = SceneConfig(n_objects=50, seed="determinism_test")

        scene1 = generate_scene(config)
        scene2 = generate_scene(config)

        hash1 = compute_scene_hash(scene1)
        hash2 = compute_scene_hash(scene2)

        assert hash1 == hash2, "Scene hash should be identical for same config"

        # Verify object positions match
        for obj1, obj2 in zip(scene1, scene2):
            assert obj1.name == obj2.name
            assert obj1.x == obj2.x
            assert obj1.y == obj2.y
            assert obj1.z == obj2.z

    def test_different_seeds_different_scenes(self):
        """Different seeds produce different scenes."""
        config1 = SceneConfig(n_objects=50, seed="seed_a")
        config2 = SceneConfig(n_objects=50, seed="seed_b")

        scene1 = generate_scene(config1)
        scene2 = generate_scene(config2)

        hash1 = compute_scene_hash(scene1)
        hash2 = compute_scene_hash(scene2)

        assert hash1 != hash2, "Different seeds should produce different scenes"

    def test_manager_scene_determinism(self):
        """Manager creates deterministic scenes."""
        mgr1 = SpatialMemoryManager()
        mgr2 = SpatialMemoryManager()

        scene1 = mgr1.create_scene("test", n_objects=100, seed="manager_test")
        scene2 = mgr2.create_scene("test", n_objects=100, seed="manager_test")

        assert scene1.scene_hash == scene2.scene_hash


# ============================================================================
# Tier 1: Query Accuracy Tests (vs Brute Force Ground Truth)
# ============================================================================

class TestQueryAccuracy:
    """Verify query results match brute-force ground truth."""

    @pytest.fixture
    def scene_and_backend(self):
        """Create a test scene with known properties."""
        config = SceneConfig(n_objects=100, seed="accuracy_test", distribution="uniform")
        objects = generate_scene(config)
        backend = FibPi3DBackend(objects, use_original_coords=True)
        return objects, backend

    def test_distance_accuracy(self, scene_and_backend):
        """Distance query matches numpy computation."""
        objects, backend = scene_and_backend

        for i in range(min(10, len(objects))):
            for j in range(i + 1, min(10, len(objects))):
                obj_a = objects[i]
                obj_b = objects[j]

                # Backend result
                backend_dist = backend.distance(obj_a.name, obj_b.name)

                # Ground truth (numpy)
                ground_truth = float(np.linalg.norm(obj_a.pos - obj_b.pos))

                assert abs(backend_dist - ground_truth) < 1e-10, \
                    f"Distance mismatch: {backend_dist} vs {ground_truth}"

    def test_near_accuracy(self, scene_and_backend):
        """Near query is consistent with distance."""
        objects, backend = scene_and_backend

        for i in range(min(10, len(objects))):
            for j in range(i + 1, min(10, len(objects))):
                obj_a = objects[i]
                obj_b = objects[j]
                dist = backend.distance(obj_a.name, obj_b.name)

                # Test with threshold slightly above and below actual distance
                assert backend.near(obj_a.name, obj_b.name, dist + 0.01) is True
                assert backend.near(obj_a.name, obj_b.name, dist - 0.01) is False

    def test_k_nearest_accuracy(self, scene_and_backend):
        """k_nearest returns correct neighbors in order."""
        objects, backend = scene_and_backend
        target = objects[0]
        k = 5

        # Backend result
        neighbors = backend.k_nearest(target.name, k)

        # Ground truth: compute all distances and sort
        distances = []
        for obj in objects[1:]:
            dist = float(np.linalg.norm(target.pos - obj.pos))
            distances.append((obj.name, dist))
        distances.sort(key=lambda x: x[1])
        ground_truth = distances[:k]

        # Verify same neighbors in same order
        assert len(neighbors) == len(ground_truth)
        for (n_name, n_dist), (gt_name, gt_dist) in zip(neighbors, ground_truth):
            assert n_name == gt_name, f"Neighbor mismatch: {n_name} vs {gt_name}"
            assert abs(n_dist - gt_dist) < 1e-10, f"Distance mismatch"

    def test_between_accuracy(self, scene_and_backend):
        """Between query uses correct geometric predicate."""
        objects, backend = scene_and_backend

        # Create a known between scenario
        # Place obj_b on the line between obj_a and obj_c
        a = objects[0].pos
        c = objects[50].pos
        midpoint = (a + c) / 2

        # Find object closest to midpoint
        closest_to_mid = None
        min_dist = float('inf')
        for obj in objects[1:50]:
            dist = np.linalg.norm(obj.pos - midpoint)
            if dist < min_dist:
                min_dist = dist
                closest_to_mid = obj

        if closest_to_mid and min_dist < 5.0:  # Reasonable tolerance
            # This object should be "between" with sufficient tolerance
            result = backend.between(
                objects[0].name,
                closest_to_mid.name,
                objects[50].name,
                tolerance=min_dist + 0.1
            )
            # Note: may or may not be between depending on projection

    def test_left_of_above_accuracy(self, scene_and_backend):
        """Directional queries use correct coordinate comparison."""
        objects, backend = scene_and_backend

        for i in range(min(20, len(objects))):
            for j in range(i + 1, min(20, len(objects))):
                obj_a = objects[i]
                obj_b = objects[j]

                # left_of: a.x < b.x
                expected_left = obj_a.x < obj_b.x
                assert backend.left_of(obj_a.name, obj_b.name) == expected_left

                # above: a.z > b.z
                expected_above = obj_a.z > obj_b.z
                assert backend.above(obj_a.name, obj_b.name) == expected_above


# ============================================================================
# Tier 1: Manager Integration Tests
# ============================================================================

class TestManagerIntegration:
    """Test SpatialMemoryManager functionality."""

    def test_create_and_get_scene(self):
        """Can create and retrieve scenes."""
        mgr = SpatialMemoryManager()

        scene = mgr.create_scene("test_scene", n_objects=50)
        assert scene is not None
        assert scene.scene_id == "test_scene"
        assert len(scene.objects) == 50

        # Can retrieve
        retrieved = mgr.get_scene("test_scene")
        assert retrieved is not None
        assert retrieved.scene_hash == scene.scene_hash

    def test_get_or_create_scene(self):
        """get_or_create returns existing or creates new."""
        mgr = SpatialMemoryManager()

        # First call creates
        scene1 = mgr.get_or_create_scene("lazy_scene", n_objects=30)
        assert scene1 is not None

        # Second call returns existing
        scene2 = mgr.get_or_create_scene("lazy_scene", n_objects=100)  # Different config ignored
        assert scene2.scene_hash == scene1.scene_hash

    def test_list_scenes(self):
        """Can list all scenes."""
        mgr = SpatialMemoryManager()

        mgr.create_scene("scene_a")
        mgr.create_scene("scene_b")
        mgr.create_scene("scene_c")

        scenes = mgr.list_scenes()
        assert len(scenes) == 3
        assert "scene_a" in scenes
        assert "scene_b" in scenes
        assert "scene_c" in scenes

    def test_delete_scene(self):
        """Can delete scenes."""
        mgr = SpatialMemoryManager()

        mgr.create_scene("to_delete")
        assert mgr.get_scene("to_delete") is not None

        result = mgr.delete_scene("to_delete")
        assert result is True
        assert mgr.get_scene("to_delete") is None

        # Delete non-existent returns False
        result = mgr.delete_scene("nonexistent")
        assert result is False

    def test_query_via_manager(self):
        """Can execute queries via manager."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("query_test", n_objects=50)

        # k_nearest query
        result = mgr.query("query_test", "k_nearest", "obj_0000", k=3)
        assert "error" not in result or result.get("error") is None
        assert len(result["result"]) == 3
        assert result["latency_us"] > 0

        # near query
        result = mgr.query("query_test", "near", "obj_0000", b="obj_0001", threshold=100.0)
        assert "error" not in result or result.get("error") is None
        assert isinstance(result["result"], bool)

        # distance query
        result = mgr.query("query_test", "distance", "obj_0000", b="obj_0001")
        assert "error" not in result or result.get("error") is None
        assert isinstance(result["result"], float)

    def test_batch_query(self):
        """Can execute batch queries."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("batch_test", n_objects=50)

        queries = [
            {"query_type": "k_nearest", "a": "obj_0000", "k": 5},
            {"query_type": "near", "a": "obj_0000", "b": "obj_0001", "threshold": 10.0},
            {"query_type": "distance", "a": "obj_0000", "b": "obj_0025"},
            {"query_type": "left_of", "a": "obj_0000", "b": "obj_0010"},
            {"query_type": "above", "a": "obj_0005", "b": "obj_0015"},
        ]

        results = mgr.batch_query("batch_test", queries)
        assert len(results) == 5

        # All should succeed
        for r in results:
            assert r.get("error") is None, f"Query failed: {r}"

    def test_query_missing_scene(self):
        """Query on missing scene returns error."""
        mgr = SpatialMemoryManager()

        result = mgr.query("nonexistent", "k_nearest", "obj_0000")
        assert "error" in result
        assert "not found" in result["error"].lower()


# ============================================================================
# Tier 1: Distribution Tests
# ============================================================================

class TestDistributions:
    """Test different scene distributions."""

    def test_uniform_distribution(self):
        """Uniform distribution spreads objects evenly."""
        config = SceneConfig(n_objects=1000, distribution="uniform", bounds=(-10.0, 10.0))
        objects = generate_scene(config)

        xs = [obj.x for obj in objects]
        ys = [obj.y for obj in objects]
        zs = [obj.z for obj in objects]

        # Check spread covers the range
        assert min(xs) < -8.0 and max(xs) > 8.0
        assert min(ys) < -8.0 and max(ys) > 8.0
        assert min(zs) < -8.0 and max(zs) > 8.0

    def test_gaussian_distribution(self):
        """Gaussian distribution clusters around center."""
        config = SceneConfig(n_objects=1000, distribution="gaussian", bounds=(-10.0, 10.0))
        objects = generate_scene(config)

        xs = [obj.x for obj in objects]
        mean_x = sum(xs) / len(xs)

        # Mean should be near center (0)
        assert abs(mean_x) < 1.0, f"Gaussian mean should be near 0, got {mean_x}"

    def test_clustered_distribution(self):
        """Clustered distribution creates distinct groups."""
        config = SceneConfig(n_objects=500, distribution="clustered", bounds=(-10.0, 10.0))
        objects = generate_scene(config)

        # Objects with same cluster_id (i % 5) should be closer together
        # This is a weak test but verifies the clustering logic runs
        assert len(objects) == 500


# ============================================================================
# Tier 1: Performance Benchmarks
# ============================================================================

class TestPerformance:
    """Basic performance verification."""

    def test_k_nearest_performance(self):
        """k_nearest on 1000 objects completes quickly."""
        import time

        mgr = SpatialMemoryManager()
        mgr.create_scene("perf_test", n_objects=1000)

        start = time.perf_counter()
        for i in range(100):
            mgr.query("perf_test", "k_nearest", f"obj_{i:04d}", k=10)
        elapsed = time.perf_counter() - start

        # 100 queries on 1000 objects should complete in < 1 second
        assert elapsed < 1.0, f"k_nearest too slow: {elapsed:.3f}s for 100 queries"

    def test_batch_query_performance(self):
        """Batch queries are efficient."""
        import time

        mgr = SpatialMemoryManager()
        mgr.create_scene("batch_perf", n_objects=500)

        queries = [
            {"query_type": "k_nearest", "a": f"obj_{i:04d}", "k": 5}
            for i in range(30)
        ]

        start = time.perf_counter()
        results = mgr.batch_query("batch_perf", queries)
        elapsed = time.perf_counter() - start

        assert len(results) == 30
        assert elapsed < 0.5, f"Batch query too slow: {elapsed:.3f}s"


# ============================================================================
# Tier 1: Edge Cases
# ============================================================================

class TestEdgeCases:
    """Test edge cases and error handling."""

    def test_query_nonexistent_object(self):
        """Query with non-existent object name handles gracefully."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("edge_test", n_objects=10)

        # k_nearest returns empty for non-existent target
        result = mgr.query("edge_test", "k_nearest", "nonexistent_obj")
        assert result["result"] == []

        # near returns False for non-existent objects
        result = mgr.query("edge_test", "near", "nonexistent", b="obj_0000", threshold=1.0)
        assert result["result"] is False

        # distance returns inf for non-existent objects
        result = mgr.query("edge_test", "distance", "nonexistent", b="obj_0000")
        assert result["result"] == float('inf')

    def test_missing_required_params(self):
        """Queries with missing params return error."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("param_test", n_objects=10)

        # near without b
        result = mgr.query("param_test", "near", "obj_0000")
        assert "error" in result

        # between without c
        result = mgr.query("param_test", "between", "obj_0000", b="obj_0001")
        assert "error" in result

    def test_small_scene(self):
        """Very small scenes work correctly."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("tiny", n_objects=2)

        result = mgr.query("tiny", "k_nearest", "obj_0000", k=5)
        # Only 1 other object, so should return 1
        assert len(result["result"]) == 1

        result = mgr.query("tiny", "distance", "obj_0000", b="obj_0001")
        assert result["result"] > 0

    def test_unknown_query_type(self):
        """Unknown query type returns error."""
        mgr = SpatialMemoryManager()
        mgr.create_scene("unknown_test", n_objects=10)

        result = mgr.query("unknown_test", "invalid_query_type", "obj_0000")
        assert "error" in result
        assert "unknown" in result["error"].lower()


# ============================================================================
# Tier 2: Boolean Confusion Matrix Tests (WO-SPATIAL-METRICS-HARDEN-01)
# ============================================================================

class TestBooleanConfusionMetrics:
    """Test confusion matrix computation for boolean queries."""

    def test_confusion_matrix_computation(self):
        """Verify confusion matrix counts are correct."""
        from tools.svs_benchmark.query_suite import (
            compute_boolean_confusion, QueryResult, QueryType
        )

        # Simulate: ground truth has 10 True, 90 False
        # Test system finds 8 True, misses 2, has 5 false alarms
        results = []
        truth = []

        # 8 TP (both True)
        for i in range(8):
            results.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"a{i}"}, result=True,
                latency_us=1.0, backend="test"
            ))
            truth.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"a{i}"}, result=True,
                latency_us=1.0, backend="ground_truth"
            ))

        # 2 FN (truth=True, test=False) - MISSED
        for i in range(2):
            results.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"b{i}"}, result=False,
                latency_us=1.0, backend="test"
            ))
            truth.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"b{i}"}, result=True,
                latency_us=1.0, backend="ground_truth"
            ))

        # 5 FP (truth=False, test=True) - FALSE ALARM
        for i in range(5):
            results.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"c{i}"}, result=True,
                latency_us=1.0, backend="test"
            ))
            truth.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"c{i}"}, result=False,
                latency_us=1.0, backend="ground_truth"
            ))

        # 85 TN (both False)
        for i in range(85):
            results.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"d{i}"}, result=False,
                latency_us=1.0, backend="test"
            ))
            truth.append(QueryResult(
                query_type=QueryType.NEAR, args={"a": f"d{i}"}, result=False,
                latency_us=1.0, backend="ground_truth"
            ))

        confusion = compute_boolean_confusion(results, truth)

        # Verify counts
        assert confusion["confusion_matrix"]["TP"] == 8
        assert confusion["confusion_matrix"]["FN"] == 2
        assert confusion["confusion_matrix"]["FP"] == 5
        assert confusion["confusion_matrix"]["TN"] == 85
        assert confusion["total_boolean_queries"] == 100

        # Verify derived metrics
        # Recall = TP / (TP + FN) = 8 / 10 = 0.8
        assert abs(confusion["recall"] - 0.8) < 1e-6

        # Precision = TP / (TP + FP) = 8 / 13
        assert abs(confusion["precision"] - 8/13) < 1e-6

        # Base rate = (TP + FN) / total = 10 / 100 = 0.1
        assert abs(confusion["base_rate_true"] - 0.1) < 1e-4

        # Trivial accuracy = max(0.1, 0.9) = 0.9
        assert abs(confusion["trivial_accuracy"] - 0.9) < 1e-4

    def test_confusion_empty_batch(self):
        """Confusion matrix handles empty batches gracefully."""
        from tools.svs_benchmark.query_suite import (
            compute_boolean_confusion, QueryResult, QueryType
        )

        # Non-boolean results only (k_nearest returns lists)
        results = [
            QueryResult(
                query_type=QueryType.K_NEAREST, args={}, result=[("obj1", 1.0)],
                latency_us=1.0, backend="test"
            )
        ]
        truth = [
            QueryResult(
                query_type=QueryType.K_NEAREST, args={}, result=[("obj1", 1.0)],
                latency_us=1.0, backend="ground_truth"
            )
        ]

        confusion = compute_boolean_confusion(results, truth)
        assert "error" in confusion

    def test_recall_reveals_missed_relations(self):
        """Recall metric reveals when system misses true relations."""
        from tools.svs_benchmark.query_suite import (
            compute_boolean_confusion, QueryResult, QueryType
        )

        # Scenario: SVS graph misses 50% of true "between" relations
        # because edges don't exist

        results = []
        truth = []

        # Ground truth: 20 actual between relations
        # System finds only 10 (50% recall)
        for i in range(10):
            # Found (TP)
            results.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=True,
                latency_us=1.0, backend="svs"
            ))
            truth.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=True,
                latency_us=1.0, backend="brute"
            ))

        for i in range(10):
            # Missed (FN)
            results.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=False,
                latency_us=1.0, backend="svs"
            ))
            truth.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=True,
                latency_us=1.0, backend="brute"
            ))

        # 80 true negatives
        for i in range(80):
            results.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=False,
                latency_us=1.0, backend="svs"
            ))
            truth.append(QueryResult(
                query_type=QueryType.BETWEEN, args={}, result=False,
                latency_us=1.0, backend="brute"
            ))

        confusion = compute_boolean_confusion(results, truth)

        # Exact match rate is high (90%) but misleading
        # because base rate is low (20%)
        exact_match = (10 + 80) / 100  # 90%

        # Recall reveals the problem: only 50% of true relations found
        assert abs(confusion["recall"] - 0.5) < 1e-6

        # Base rate shows why high accuracy is trivial
        assert abs(confusion["base_rate_true"] - 0.2) < 1e-4


class TestBalancedQuerySampling:
    """Test balanced query generation."""

    def test_balanced_near_queries_target_rate(self):
        """Balanced near queries achieve approximately target true rate."""
        from tools.svs_benchmark.generators import generate_balanced_near_queries
        from tools.svs_benchmark.scene_generator import generate_scene

        scene = generate_scene(200, "balanced_test")
        queries = generate_balanced_near_queries(
            scene, n_queries=100, threshold=5.0, target_true_rate=0.4
        )

        assert len(queries) == 100

        # Verify query structure
        for q in queries:
            assert "a_name" in q
            assert "b_name" in q
            assert "threshold" in q
            assert q["a_name"] != q["b_name"]

    def test_balanced_between_queries_target_rate(self):
        """Balanced between queries achieve approximately target true rate."""
        from tools.svs_benchmark.generators import generate_balanced_between_queries
        from tools.svs_benchmark.scene_generator import generate_scene

        scene = generate_scene(200, "balanced_between_test")
        queries = generate_balanced_between_queries(
            scene, n_queries=100, tolerance=5.0, target_true_rate=0.4
        )

        # May get fewer queries if not enough candidates
        assert len(queries) <= 100

        # Verify query structure
        for q in queries:
            assert "a_name" in q
            assert "b_name" in q
            assert "c_name" in q
            assert "tolerance" in q
            # All three must be distinct
            assert len({q["a_name"], q["b_name"], q["c_name"]}) == 3

    def test_estimate_true_rate(self):
        """True rate estimation provides confidence interval."""
        from tools.svs_benchmark.generators import estimate_true_rate
        from tools.svs_benchmark.scene_generator import generate_scene

        scene = generate_scene(100, "estimate_test")
        est = estimate_true_rate(scene, "near", threshold=5.0, n_samples=200)

        assert "true_rate" in est
        assert "ci_lower" in est
        assert "ci_upper" in est
        assert 0 <= est["true_rate"] <= 1
        assert est["ci_lower"] <= est["true_rate"] <= est["ci_upper"]

    def test_balanced_queries_deterministic(self):
        """Balanced query generation is deterministic."""
        from tools.svs_benchmark.generators import generate_balanced_near_queries
        from tools.svs_benchmark.scene_generator import generate_scene

        scene = generate_scene(100, "determinism_test")

        queries1 = generate_balanced_near_queries(
            scene, n_queries=50, threshold=5.0, seed="fixed_seed"
        )
        queries2 = generate_balanced_near_queries(
            scene, n_queries=50, threshold=5.0, seed="fixed_seed"
        )

        # Queries should be identical
        assert len(queries1) == len(queries2)
        for q1, q2 in zip(queries1, queries2):
            assert q1 == q2


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
