"""Benchmark QKMeans against sklearn k-means++ seeding."""

import time
import numpy as np
from sklearn.cluster import kmeans_plusplus
from sklearn.datasets import make_blobs

try:
    from qkmeans import QKMeansSeeding, get_simd_level
    QKMEANS_AVAILABLE = True
except ImportError:
    QKMEANS_AVAILABLE = False
    print("Warning: qkmeans not installed, skipping QKMeans benchmarks")


def benchmark_sklearn_kmeanspp(X, k, n_runs=5):
    """Benchmark sklearn's k-means++ implementation."""
    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        centers, indices = kmeans_plusplus(X, n_clusters=k, random_state=None)
        elapsed = (time.perf_counter() - start) * 1000  # ms
        times.append(elapsed)
    return np.mean(times), np.std(times)


def benchmark_qkmeans(X, k, anns='lsh', n_runs=5, **kwargs):
    """Benchmark QKMeans implementation."""
    if not QKMEANS_AVAILABLE:
        return float('nan'), float('nan')

    times = []
    for i in range(n_runs):
        seeder = QKMeansSeeding(n_clusters=k, anns=anns, random_state=i, **kwargs)
        start = time.perf_counter()
        seeder.fit(X)
        elapsed = (time.perf_counter() - start) * 1000  # ms
        times.append(elapsed)
    return np.mean(times), np.std(times)


def run_benchmark(n_samples, n_features, n_clusters, n_runs=5):
    """Run benchmark for given parameters."""
    print(f"\n{'='*60}")
    print(f"Benchmark: n={n_samples}, d={n_features}, k={n_clusters}")
    print(f"{'='*60}")

    # Generate data
    print("Generating data...", end=" ", flush=True)
    X, _ = make_blobs(
        n_samples=n_samples,
        n_features=n_features,
        centers=min(n_clusters, 50),  # Limit actual centers for make_blobs
        random_state=42
    )
    X = X.astype(np.float64)
    print("done")

    results = []

    # Sklearn k-means++
    print("Running sklearn k-means++...", end=" ", flush=True)
    mean_time, std_time = benchmark_sklearn_kmeanspp(X, n_clusters, n_runs)
    print(f"{mean_time:.2f} ± {std_time:.2f} ms")
    results.append(("sklearn", mean_time, std_time))

    if QKMEANS_AVAILABLE:
        # QKMeans with LSH
        print("Running QKMeans (LSH)...", end=" ", flush=True)
        mean_time, std_time = benchmark_qkmeans(X, n_clusters, 'lsh', n_runs)
        print(f"{mean_time:.2f} ± {std_time:.2f} ms")
        results.append(("QKMeans-LSH", mean_time, std_time))

        # QKMeans with brute force (for comparison)
        if n_clusters <= 100:  # Only for small k
            print("Running QKMeans (brute force)...", end=" ", flush=True)
            mean_time, std_time = benchmark_qkmeans(X, n_clusters, 'brute_force', n_runs)
            print(f"{mean_time:.2f} ± {std_time:.2f} ms")
            results.append(("QKMeans-BF", mean_time, std_time))

    # Print summary
    print("\nSummary:")
    print("-" * 40)
    sklearn_time = results[0][1]
    for name, mean_time, std_time in results:
        speedup = sklearn_time / mean_time if mean_time > 0 else float('nan')
        print(f"{name:15s}: {mean_time:8.2f} ms  (speedup: {speedup:.2f}x)")

    return results


def main():
    """Run benchmarks."""
    print("QKMeans Benchmark Suite")
    print("=" * 60)

    if QKMEANS_AVAILABLE:
        print(f"SIMD Level: {get_simd_level()}")
    print(f"NumPy threads: {np.__config__.show()}" if hasattr(np.__config__, 'show') else "")

    # Benchmark configurations
    configs = [
        # (n_samples, n_features, n_clusters)
        (1000, 10, 10),
        (10000, 50, 50),
        (10000, 128, 100),
        (50000, 128, 200),
        (100000, 128, 500),
    ]

    all_results = {}
    for n, d, k in configs:
        results = run_benchmark(n, d, k, n_runs=3)
        all_results[(n, d, k)] = results

    # Final summary table
    print("\n" + "=" * 60)
    print("Final Summary")
    print("=" * 60)
    print(f"{'n':>8s} {'d':>5s} {'k':>5s} {'sklearn':>12s} {'QKMeans-LSH':>12s} {'Speedup':>8s}")
    print("-" * 60)

    for (n, d, k), results in all_results.items():
        sklearn_time = results[0][1]
        qkmeans_time = results[1][1] if len(results) > 1 else float('nan')
        speedup = sklearn_time / qkmeans_time if qkmeans_time > 0 else float('nan')
        print(f"{n:8d} {d:5d} {k:5d} {sklearn_time:10.2f}ms {qkmeans_time:10.2f}ms {speedup:7.2f}x")


if __name__ == "__main__":
    main()
