import numpy as np
import time
from scipy import stats
import pandas as pd

def benchmark_regression(n, n_features=10, n_trials=50):
    """Benchmark OLS regression across different sample sizes."""
    times = []

    for _ in range(n_trials):
        X = np.random.randn(n, n_features)
        beta = np.random.randn(n_features)
        y = X @ beta + np.random.randn(n) * 0.5

        start = time.perf_counter()
        # OLS via normal equations: β = (X'X)^(-1) X'y
        XtX = X.T @ X
        Xty = X.T @ y
        beta_hat = np.linalg.solve(XtX, Xty)
        residuals = y - X @ beta_hat
        se = np.sqrt(np.diag(np.linalg.inv(XtX) * np.var(residuals)))
        t_vals = beta_hat / se
        p_vals = 2 * (1 - stats.t.cdf(np.abs(t_vals), df=n - n_features))
        elapsed = time.perf_counter() - start
        times.append(elapsed)

    return {
        "n": n,
        "mean_time_ms": np.mean(times) * 1000,
        "std_time_ms": np.std(times) * 1000,
        "median_time_ms": np.median(times) * 1000
    }

# Benchmark across sample sizes
sample_sizes = [100, 500, 1000, 5000, 10000, 50000]
results = []

print("OLS Regression Benchmark (NumPy)")
print("=" * 55)
print(f"{'N':>8} {'Mean (ms)':>12} {'Std (ms)':>12} {'Median (ms)':>12}")
print("-" * 55)

for n in sample_sizes:
    result = benchmark_regression(n)
    results.append(result)
    print(f"{result['n']:>8} {result['mean_time_ms']:>12.3f} "
          f"{result['std_time_ms']:>12.3f} {result['median_time_ms']:>12.3f}")

# Scaling analysis
df_results = pd.DataFrame(results)
# OLS should scale roughly O(n * p^2) for the solve step
df_results["time_per_obs"] = df_results["mean_time_ms"] / df_results["n"]

print("\nScaling Analysis:")
print(df_results[["n", "time_per_obs"]].to_string(index=False))

# Comparison: scipy.stats vs manual implementation
np.random.seed(42)
X = np.random.randn(1000, 5)
y = 2.0 * X[:, 0] + 1.5 * X[:, 1] + np.random.randn(1000) * 0.5

# Manual
start = time.perf_counter()
XtX = X.T @ X
Xty = X.T @ y
beta = np.linalg.solve(XtX, Xty)
res = y - X @ beta
se = np.sqrt(np.diag(np.linalg.inv(XtX) * np.var(res)))
t_manual = time.perf_counter() - start

# statsmodels
import statsmodels.api as sm
X_sm = sm.add_constant(X)
start = time.perf_counter()
model = sm.OLS(y, X_sm).fit()
t_sm = time.perf_counter() - start

print(f"\nManual OLS: {t_manual*1000:.3f} ms")
print(f"statsmodels OLS: {t_sm*1000:.3f} ms")
print(f"Ratio (sm/manual): {t_sm/t_manual:.1f}x")