test_basic.py

"""
Comprehensive test for bbox pruning AND balanced tree construction
Tests BOTH the bug fix and the Phase C optimization
"""

from spatiox import _spatio_core as sc
import numpy as np


def test_bbox_pruning_with_balanced_tree():
    """Verify bbox pruning works with balanced tree construction"""
    
    print("="*80)
    print("COMPLETE TEST: Bbox Pruning + Balanced Tree Construction")
    print("="*80)
    print()
    
    # Create dataset with clear spatial clustering
    idx = sc.SpatioIndexCore()
    
    # Cluster 1: San Francisco area
    np.random.seed(42)
    cluster1_lat = 37.7749 + np.random.randn(1000) * 0.02
    cluster1_lon = -122.4194 + np.random.randn(1000) * 0.02
    cluster1_time = np.linspace(1000, 1500, 1000)
    
    # Cluster 2: Los Angeles area (far away ~560km)
    cluster2_lat = 34.0522 + np.random.randn(1000) * 0.02
    cluster2_lon = -118.2437 + np.random.randn(1000) * 0.02
    cluster2_time = np.linspace(1500, 2000, 1000)
    
    # Combine
    data = (list(zip(cluster1_lat, cluster1_lon, cluster1_time)) +
            list(zip(cluster2_lat, cluster2_lon, cluster2_time)))
    
    print(f"Dataset: 2,000 records in 2 distant clusters")
    print(f"  Cluster 1: San Francisco (~1,000 points)")
    print(f"  Cluster 2: Los Angeles (~1,000 points, ~560km away)")
    print()
    
    # Test 1: WITHOUT build() (incremental insertion - Phase A)
    print("-" * 80)
    print("TEST 1: Phase A - Incremental Insertion (no build)")
    print("-" * 80)
    
    idx.bulk_insert(data)
    # NO build() call - tree is unbalanced
    
    sf_lat, sf_lon = 37.7749, -122.4194
    results_phase_a, stats_phase_a = idx.query_radius_time_instrumented(
        sf_lat, sf_lon, 5.0, 1000, 2000
    )
    
    total_nodes_phase_a = idx.get_index_stats().spatial_nodes
    visit_rate_phase_a = (stats_phase_a.spatial_nodes_visited / total_nodes_phase_a) * 100
    
    print(f"Results: {len(results_phase_a)} records found")
    print()
    print(f"Performance (Phase A - unbalanced):")
    print(f"  Nodes visited: {stats_phase_a.spatial_nodes_visited}/{total_nodes_phase_a} ({visit_rate_phase_a:.1f}%)")
    print(f"  Distance checks: {stats_phase_a.spatial_distance_checks}")
    print(f"  Bbox prunes: {stats_phase_a.spatial_bbox_prunes}")
    print(f"  Distance prunes: {stats_phase_a.spatial_distance_prunes}")
    print()
    
    # Check if bbox pruning is working
    if stats_phase_a.spatial_bbox_prunes > 0:
        prune_rate = (stats_phase_a.spatial_bbox_prunes / stats_phase_a.spatial_nodes_visited) * 100
        print(f"✅ Bbox pruning IS working: {stats_phase_a.spatial_bbox_prunes} subtrees pruned ({prune_rate:.1f}%)")
    else:
        print(f"❌ Bbox pruning NOT working: 0 subtrees pruned")
        print(f"   BUG: Bounding boxes not being checked!")
        return False
    
    print()
    
    # Test 2: WITH build() (balanced tree - Phase C)
    print("-" * 80)
    print("TEST 2: Phase C - Balanced Tree Construction")
    print("-" * 80)
    
    print("Calling build() to construct balanced KD-tree...")
    idx.build()
    print("✓ Tree rebuilt with median-based splits")
    print()
    
    results_phase_c, stats_phase_c = idx.query_radius_time_instrumented(
        sf_lat, sf_lon, 5.0, 1000, 2000
    )
    
    total_nodes_phase_c = idx.get_index_stats().spatial_nodes
    visit_rate_phase_c = (stats_phase_c.spatial_nodes_visited / total_nodes_phase_c) * 100
    
    print(f"Results: {len(results_phase_c)} records found")
    print()
    print(f"Performance (Phase C - balanced):")
    print(f"  Nodes visited: {stats_phase_c.spatial_nodes_visited}/{total_nodes_phase_c} ({visit_rate_phase_c:.1f}%)")
    print(f"  Distance checks: {stats_phase_c.spatial_distance_checks}")
    print(f"  Bbox prunes: {stats_phase_c.spatial_bbox_prunes}")
    print(f"  Distance prunes: {stats_phase_c.spatial_distance_prunes}")
    print()
    
    # Calculate improvement
    nodes_improvement = ((stats_phase_a.spatial_nodes_visited - stats_phase_c.spatial_nodes_visited) / 
                         stats_phase_a.spatial_nodes_visited * 100)
    dist_improvement = ((stats_phase_a.spatial_distance_checks - stats_phase_c.spatial_distance_checks) /
                        stats_phase_a.spatial_distance_checks * 100)
    
    print(f"📊 Improvement from balanced tree:")
    print(f"  Nodes visited reduced: {nodes_improvement:.1f}%")
    print(f"  Distance checks reduced: {dist_improvement:.1f}%")
    print()
    
    # Verify correctness
    print("-" * 80)
    print("CORRECTNESS VERIFICATION")
    print("-" * 80)
    
    if set(results_phase_a) == set(results_phase_c):
        print("✅ PASS: Results identical (balanced tree is exact)")
        print(f"   Both returned {len(results_phase_a)} records")
    else:
        print("❌ FAIL: Results differ!")
        print(f"   Phase A: {len(results_phase_a)} records")
        print(f"   Phase C: {len(results_phase_c)} records")
        return False
    
    print()
    
    # Summary
    print("=" * 80)
    print("SUMMARY")
    print("=" * 80)
    
    success = True
    
    if stats_phase_a.spatial_bbox_prunes > 0:
        print("✅ Bbox pruning is WORKING")
    else:
        print("❌ Bbox pruning FAILED")
        success = False
    
    if stats_phase_c.spatial_bbox_prunes > stats_phase_a.spatial_bbox_prunes:
        print(f"✅ Balanced tree improves pruning: {stats_phase_c.spatial_bbox_prunes} vs {stats_phase_a.spatial_bbox_prunes} prunes")
    else:
        print(f"⚠️  Balanced tree pruning: {stats_phase_c.spatial_bbox_prunes} vs {stats_phase_a.spatial_bbox_prunes} prunes")
    
    if nodes_improvement > 10:
        print(f"✅ Balanced tree reduces nodes visited: {nodes_improvement:.1f}% improvement")
    else:
        print(f"⚠️  Balanced tree node reduction: {nodes_improvement:.1f}% (expected >10%)")
    
    if set(results_phase_a) == set(results_phase_c):
        print("✅ Correctness maintained (exact results)")
    else:
        print("❌ Correctness violated")
        success = False
    
    print()
    
    if success:
        print("🎉 ALL TESTS PASSED - Bbox pruning AND balanced tree working!")
        return True
    else:
        print("⚠️  SOME TESTS FAILED - See details above")
        return False


def test_query_far_from_data():
    """Test query far from all data - should prune aggressively"""
    
    print("\n" + "=" * 80)
    print("TEST: Query Far from Data (Maximum Pruning)")
    print("=" * 80)
    print()
    
    idx = sc.SpatioIndexCore()
    
    # Data in San Francisco
    np.random.seed(123)
    data = [(37.7749 + np.random.randn()*0.02, -122.4194 + np.random.randn()*0.02, 1000.0 + i)
            for i in range(1000)]
    
    idx.bulk_insert(data)
    idx.build()
    
    # Query in Nevada (far from data)
    nevada_lat, nevada_lon = 39.0, -117.0
    
    results, stats = idx.query_radius_time_instrumented(
        nevada_lat, nevada_lon, 50.0, 0, 2000
    )
    
    total_nodes = idx.get_index_stats().spatial_nodes
    visit_rate = (stats.spatial_nodes_visited / total_nodes) * 100
    
    print(f"Query: 50km radius in Nevada (far from all data)")
    print(f"Results: {len(results)} records")
    print()
    print(f"Performance:")
    print(f"  Nodes visited: {stats.spatial_nodes_visited}/{total_nodes} ({visit_rate:.1f}%)")
    print(f"  Distance checks: {stats.spatial_distance_checks}")
    print(f"  Bbox prunes: {stats.spatial_bbox_prunes}")
    print()
    
    if visit_rate < 30:
        print(f"✅ PASS: Visited only {visit_rate:.1f}% of tree (excellent pruning)")
    else:
        print(f"⚠️  Visited {visit_rate:.1f}% of tree (expected <30%)")
    
    if stats.spatial_bbox_prunes > stats.spatial_distance_checks:
        print(f"✅ PASS: More bbox prunes than distance checks (efficient)")
    else:
        print(f"⚠️  More distance checks than bbox prunes (less efficient)")
    
    print()


def test_small_radius_efficiency():
    """Test small radius query - should prune aggressively"""
    
    print("=" * 80)
    print("TEST: Small Radius Query (Aggressive Pruning)")
    print("=" * 80)
    print()
    
    idx = sc.SpatioIndexCore()
    
    # Spread out data
    np.random.seed(456)
    data = [(37.7 + np.random.randn()*0.2, -122.4 + np.random.randn()*0.2, 1000.0 + i)
            for i in range(5000)]
    
    idx.bulk_insert(data)
    idx.build()
    
    # Very small radius
    results, stats = idx.query_radius_time_instrumented(
        37.7, -122.4, 0.5, 0, 2000
    )
    
    total_nodes = idx.get_index_stats().spatial_nodes
    visit_rate = (stats.spatial_nodes_visited / total_nodes) * 100
    prune_rate = (stats.spatial_bbox_prunes / stats.spatial_nodes_visited * 100) if stats.spatial_nodes_visited > 0 else 0
    
    print(f"Query: 0.5km radius (very small)")
    print(f"Results: {len(results)} records")
    print()
    print(f"Performance:")
    print(f"  Nodes visited: {stats.spatial_nodes_visited}/{total_nodes} ({visit_rate:.1f}%)")
    print(f"  Bbox prunes: {stats.spatial_bbox_prunes} ({prune_rate:.1f}% of visited)")
    print(f"  Distance checks: {stats.spatial_distance_checks}")
    print()
    
    if visit_rate < 40:
        print(f"✅ PASS: Small radius visited only {visit_rate:.1f}% of tree")
    else:
        print(f"⚠️  Small radius visited {visit_rate:.1f}% (expected <40%)")
    
    if prune_rate > 50:
        print(f"✅ PASS: High pruning rate {prune_rate:.1f}%")
    else:
        print(f"⚠️  Pruning rate {prune_rate:.1f}% (expected >50%)")
    
    print()


def main():
    """Run all tests"""
    
    print("\n" + "🧪" * 40)
    print("COMPREHENSIVE BBOX PRUNING + BALANCED TREE TEST SUITE")
    print("🧪" * 40 + "\n")
    
    success = test_bbox_pruning_with_balanced_tree()
    
    if success:
        test_query_far_from_data()
        test_small_radius_efficiency()
        
        print("=" * 80)
        print("✅ ALL TESTS PASSED - System is production-ready!")
        print("=" * 80)
    else:
        print("=" * 80)
        print("❌ CRITICAL FAILURES - Fix bbox pruning first")
        print("=" * 80)


if __name__ == "__main__":
    main()