import argparse
import json
import numpy as np


def math_check(d=20, gamma=0.5, lam=0.5, trials=200000, seed=0):
    # Bounded coordinate samples: x = +/- e_i, with diagonal population covariance.
    rng = np.random.default_rng(seed)
    eig = np.linspace(0.05, 1.0, d) / d
    # x_i = +/-sqrt(d*eig_i)e_i gives Sigma=diag(eig), while remaining bounded.
    idx = rng.integers(0, d, size=trials)
    signs = rng.choice([-1.0, 1.0], size=trials)
    xscale = np.sqrt(d * eig[idx])
    # P is diagonal for each sample, so estimate E[P^2] directly.
    pdiag = np.full((trials, d), 1.0 - gamma * lam)
    pdiag[np.arange(trials), idx] -= gamma * xscale**2
    ep2 = (pdiag * pdiag).mean(axis=0)
    A = 1.0 - gamma * (eig + lam)
    rhs1 = (1.0 - gamma * lam) * A
    rhs2 = (1.0 - gamma * lam) ** 2
    # Since all matrices are diagonal, these are exact eigenvalue gaps.
    gap_lemma = float(np.max(ep2 - rhs1))
    gap_contraction = float(np.max(rhs1 - rhs2))
    return {
        "max_E_P2_minus_(1-gamma-lambda)A": gap_lemma,
        "max_(1-gamma-lambda)A_minus_scalar_bound": gap_contraction,
        "gamma_lambda": gamma * lam,
        "A_min": float(A.min()),
        "A_max": float(A.max()),
        "passed": bool(gap_lemma <= 0.01 and gap_contraction <= 1e-12 and 0 <= gamma * lam <= 1),
    }


def make_data(seed, n_train=6000, n_test=3000, d=30):
    rng = np.random.default_rng(seed)
    eig = np.geomspace(1.0, 0.03, d)
    # Gaussian inputs with controlled covariance and a mildly noisy target.
    xtr = rng.normal(size=(n_train, d)) * np.sqrt(eig)
    xte = rng.normal(size=(n_test, d)) * np.sqrt(eig)
    teacher = rng.normal(size=d) / np.sqrt(np.arange(1, d + 1))
    ytr = xtr @ teacher + 0.25 * rng.normal(size=n_train)
    yte = xte @ teacher + 0.25 * rng.normal(size=n_test)
    return xtr, ytr, xte, yte


def train(x, y, xt, yt, seed, mode, gamma=0.08, lam=0.35, m=300, T=900, batch=32):
    rng = np.random.default_rng(seed + 10000)
    d = x.shape[1]
    theta = np.zeros(d)
    avg = np.zeros(d)
    avg_count = 0
    losses = []
    update_ratios = []
    for t in range(T):
        ix = rng.integers(0, len(x), size=batch)
        xb, yb = x[ix], y[ix]
        residual = xb @ theta - yb
        grad = xb.T @ residual / batch
        if mode == "constant":
            lt = lam
        elif mode == "initial":
            lt = lam if t < m else 0.0
        elif mode == "none":
            lt = 0.0
        else:
            raise ValueError(mode)
        old = theta.copy()
        theta = (1.0 - gamma * lt) * theta - gamma * grad
        update_ratios.append(float(np.linalg.norm(theta - old) / max(1.0, np.linalg.norm(old))))
        if 2 * m <= t < 3 * m:
            avg += theta
            avg_count += 1
        if (t + 1) % 100 == 0:
            losses.append(float(np.mean((x @ theta - y) ** 2) / 2))
    tail = avg / max(avg_count, 1)
    final_mse = float(np.mean((xt @ theta - yt) ** 2))
    tail_mse = float(np.mean((xt @ tail - yt) ** 2))
    return {
        "final_test_mse": final_mse,
        "tail_test_mse": tail_mse,
        "final_train_loss": losses[-1],
        "max_update_ratio": max(update_ratios),
        "loss_trace": losses,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=8)
    ap.add_argument("--out", default="results.json")
    args = ap.parse_args()
    check = math_check()
    methods = ["constant", "initial", "none"]
    all_results = {k: [] for k in methods}
    for seed in range(args.runs):
        data = make_data(seed)
        for method in methods:
            all_results[method].append(train(*data, seed, method))
    summary = {}
    for method, rows in all_results.items():
        summary[method] = {
            "final_test_mse_mean": float(np.mean([r["final_test_mse"] for r in rows])),
            "final_test_mse_std": float(np.std([r["final_test_mse"] for r in rows], ddof=1)),
            "tail_test_mse_mean": float(np.mean([r["tail_test_mse"] for r in rows])),
            "tail_test_mse_std": float(np.std([r["tail_test_mse"] for r in rows], ddof=1)),
            "final_train_loss_mean": float(np.mean([r["final_train_loss"] for r in rows])),
            "max_update_ratio_mean": float(np.mean([r["max_update_ratio"] for r in rows])),
        }
    output = {"config": {"runs": args.runs, "gamma": 0.08, "lambda": 0.35, "m": 300, "T": 900, "batch": 32}, "math_check": check, "summary": summary, "raw": all_results}
    with open(args.out, "w") as f:
        json.dump(output, f, indent=2)
    print(json.dumps({"math_check": check, "summary": summary}, indent=2))


if __name__ == "__main__":
    main()