import time
import tracemalloc
import numpy as np
import pandas as pd
from sklearn.datasets import fetch_20newsgroups_vectorized
from sklearn.decomposition import IncrementalPCA


def measure_performance(func, *args, **kwargs):
    tracemalloc.start()
    start = time.perf_counter()
    result = func(*args, **kwargs)
    elapsed = time.perf_counter() - start
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return {"time_s": elapsed, "peak_mb": peak / (1024**2)}, result


def sparse_ipca_arpack(X):
    ipca = IncrementalPCA(n_components=20, svd_solver="arpack")
    coords = ipca.fit_transform(X)
    return ipca, coords


def dense_ipca_arpack(X):
    X_dense = X.toarray()
    ipca = IncrementalPCA(n_components=20, svd_solver="arpack")
    coords = ipca.fit_transform(X_dense)
    return ipca, coords


def dense_ipca_full(X):
    X_dense = X.toarray()
    ipca = IncrementalPCA(n_components=20, svd_solver="full")
    coords = ipca.fit_transform(X_dense)
    return ipca, coords


def main():
    n_samples = 3000
    X, _ = fetch_20newsgroups_vectorized(return_X_y=True)
    X = X[:n_samples]

    methods = [
        ("Sparse ARPACK", sparse_ipca_arpack),
        ("Dense ARPACK", dense_ipca_arpack),
        ("Dense Full", dense_ipca_full),
    ]
    metrics = {}
    models = {}
    coords = {}

    print()
    print(f"\033[1mBenchmarking on {n_samples} samples...\033[0m")
    for name, func in methods:
        print(f"Running {name}...", end=" ", flush=True)
        stats, output = measure_performance(func, X)
        model, coord = output
        metrics[name] = stats
        models[name] = model
        coords[name] = coord
        print(f"Time = {stats['time_s']:.3f}s, Peak Memory = {stats['peak_mb']:.2f}MB")

    print()
    print("\033[1mVerifying results...\033[0m")
    base = "Dense Full"
    base_model = models[base]
    for name, _ in methods:
        if name == base:
            continue
        model = models[name]
        assert np.allclose(base_model.components_, model.components_)
        assert np.allclose(base_model.explained_variance_, model.explained_variance_)
        assert np.allclose(base_model.singular_values_, model.singular_values_)
        print(f"- {base} vs {name}: OK")
    print("All results are equivalent! ✅")

    print()
    print("\033[1mSummarizing performance and memory usage...\033[0m")
    base_stats = metrics[base]
    for name in methods:
        key = name[0]
        if key == base:
            metrics[key]["speedup"] = 1.0
            metrics[key]["memory_saving"] = 1.0
        else:
            t = metrics[key]["time_s"]
            m = metrics[key]["peak_mb"]
            metrics[key]["speedup"] = base_stats["time_s"] / t
            metrics[key]["memory_saving"] = base_stats["peak_mb"] / m

    df = pd.DataFrame(metrics).T
    df = df[["time_s", "peak_mb", "speedup", "memory_saving"]]
    print(df.round(3))


if __name__ == "__main__":
    main()