#!/usr/bin/env python3
"""
Benchmark script to compare NumPy's weighted_gram_matrix implementation 
with tabmat's sandwich product.
"""

import time
import numpy as np
import tabmat


def benchmark_numpy_implementation(X, weights=None, iterations=5):
    """Benchmark NumPy's weighted_gram_matrix implementation."""
    times = []
    for _ in range(iterations):
        start = time.perf_counter()
        result = np.linalg.weighted_gram_matrix(X, weights=weights)
        end = time.perf_counter()
        times.append(end - start)
    return np.mean(times), result


def benchmark_tabmat_implementation(X, weights=None, iterations=5):
    """Benchmark tabmat's sandwich product implementation."""
    # Convert to tabmat format
    if isinstance(X, np.ndarray):
        X_tabmat = tabmat.DenseMatrix(X)
    else:
        X_tabmat = X
    
    times = []
    for _ in range(iterations):
        start = time.perf_counter()
        if weights is not None:
            result = X_tabmat.sandwich(weights)
        else:
            # For no weights, use sandwich with ones
            result = X_tabmat.sandwich(np.ones(X.shape[0]))
        end = time.perf_counter()
        times.append(end - start)
    return np.mean(times), result


def run_benchmarks():
    """Run benchmarks for different matrix sizes."""
    print("Running benchmarks for weighted Gram matrix computation...")
    print("=" * 60)
    
    # Test cases with different sizes
    test_cases = [
        {"name": "Small (100x50)", "shape": (100, 50)},
        {"name": "Medium (1000x500)", "shape": (1000, 500)},
        {"name": "Large (5000x2000)", "shape": (5000, 2000)},
    ]
    
    results = []
    
    for case in test_cases:
        print(f"\nTesting {case['name']} matrix...")
        print("-" * 30)
        
        # Generate test data
        np.random.seed(42)  # For reproducibility
        X = np.random.rand(*case['shape'])
        weights = np.random.rand(case['shape'][0])
        
        # Benchmark NumPy implementation
        numpy_time, numpy_result = benchmark_numpy_implementation(X, weights)
        print(f"NumPy implementation: {numpy_time:.6f} seconds")
        
        # Benchmark tabmat implementation
        tabmat_time, tabmat_result = benchmark_tabmat_implementation(X, weights)
        print(f"Tabmat implementation: {tabmat_time:.6f} seconds")
        
        # Calculate speedup
        speedup = numpy_time / tabmat_time if tabmat_time > 0 else float('inf')
        print(f"Speedup (NumPy/tabmat): {speedup:.2f}x")
        
        # Verify results are similar (within tolerance)
        diff = np.abs(numpy_result - tabmat_result).max()
        print(f"Max difference between implementations: {diff:.2e}")
        
        results.append({
            "case": case["name"],
            "numpy_time": numpy_time,
            "tabmat_time": tabmat_time,
            "speedup": speedup,
            "max_diff": diff
        })
    
    # Print summary
    print("\n" + "=" * 60)
    print("BENCHMARK SUMMARY")
    print("=" * 60)
    print(f"{'Case':<15} {'NumPy (s)':<12} {'Tabmat (s)':<12} {'Speedup':<10} {'Max Diff':<10}")
    print("-" * 60)
    
    for result in results:
        print(f"{result['case']:<15} {result['numpy_time']:<12.6f} {result['tabmat_time']:<12.6f} "
              f"{result['speedup']:<10.2f}x {result['max_diff']:<10.2e}")


if __name__ == "__main__":
    run_benchmarks()