import json
import math
from pathlib import Path
import numpy as np

# Centered Heavy-Tail Clipping Optimizer: small reproducible numerical MVP.
# The per-example gradients are g_i = grad f(x) + heavy-tailed noise.

def clip_centered(g, c, tau):
    r = g - c
    n = np.linalg.norm(r, axis=-1, keepdims=True)
    return c + r * np.minimum(1.0, tau / (n + 1e-12))

def coordinate_median(g):
    return np.median(g, axis=0)

def grad_bias_scaling(seed=0, alpha=1.5):
    """Check E||noise 1_{||noise||>tau}|| decreases approximately tau^(1-alpha)."""
    rng = np.random.default_rng(seed)
    # scalar symmetric Pareto tails make the asserted exponent particularly clear.
    n = 2_000_000
    signs = rng.choice([-1.0, 1.0], n)
    u = rng.random(n)
    z = signs * u ** (-1.0 / alpha)  # P(|Z|>t)=t^-alpha, t>=1
    taus = np.array([2., 4., 8., 16., 32., 64.])
    vals = np.array([np.mean(np.abs(z) * (np.abs(z) > t)) for t in taus])
    slope = np.polyfit(np.log(taus), np.log(vals), 1)[0]
    # Formula C_tau is checked separately on random vectors.
    return {"taus": taus.tolist(), "tail_bias": vals.tolist(),
            "loglog_slope": float(slope), "expected_slope": float(1-alpha)}

def operator_check(seed=1):
    rng = np.random.default_rng(seed)
    g = rng.normal(size=(1000, 7)); c = rng.normal(size=7); tau = 0.73
    h = clip_centered(g, c, tau)
    residual = np.linalg.norm(h-c, axis=1)
    raw_residual = np.linalg.norm(g-c, axis=1)
    # unchanged inside ball, bounded outside, and direction preserved (up to roundoff)
    inside_err = np.max(np.abs(residual[raw_residual <= tau] - raw_residual[raw_residual <= tau]))
    max_outside = np.max(residual)
    cos = np.sum((h-c)*(g-c), axis=1) / (residual*raw_residual + 1e-30)
    return {"inside_max_error": float(inside_err), "max_output_residual": float(max_outside),
            "tau": tau, "minimum_direction_cosine": float(np.min(cos))}

def draw_gradients(x, batch, d, rng, regime):
    if regime == "gaussian":
        noise = rng.normal(0, 0.35, size=(batch, d))
        # Same occasional contamination rate in both controls, but modest magnitude.
        p, mult = 0.002, 8.
    else:
        df = 1.5
        noise = rng.standard_t(df, size=(batch, d)) * 0.20
        p, mult = 0.012, 35.
    mask = rng.random(batch) < p
    if np.any(mask):
        noise[mask] += rng.normal(size=(mask.sum(), d)) * mult
    return x[None, :] + noise

def run_one(seed, method, regime, steps=300):
    rng = np.random.default_rng(seed)
    d, batch, eta, tau = 20, 64, 0.075, 2.0
    x = np.ones(d) * 5.0
    losses = []; update_norms = []; divergent = False
    for t in range(steps):
        gs = draw_gradients(x, batch, d, rng, regime)
        if method == "sgd":
            h = gs.mean(axis=0)
        elif method == "global_clip":
            h0 = gs.mean(axis=0)
            n = np.linalg.norm(h0)
            h = h0 * min(1., tau/(n+1e-12))
        elif method == "centered":
            c = coordinate_median(gs)
            h = clip_centered(gs, c, tau).mean(axis=0)
        else:
            raise ValueError(method)
        x -= eta*h
        loss = 0.5*float(np.dot(x,x))
        losses.append(loss); update_norms.append(float(eta*np.linalg.norm(h)))
        if not np.isfinite(loss) or loss > 1e8:
            divergent = True; break
    # Catastrophic steps are unusually large parameter changes, measured uniformly.
    return {"final_loss": losses[-1], "best_loss": min(losses),
            "median_update": float(np.median(update_norms)),
            "p99_update": float(np.quantile(update_norms, .99)),
            "diverged": divergent, "losses": losses}

def aggregate():
    out = {"operator_check": operator_check(), "bias_scaling": grad_bias_scaling(), "runs": {}}
    for regime in ["heavy_tail", "gaussian"]:
        for method in ["sgd", "global_clip", "centered"]:
            rs = [run_one(s, method, regime) for s in range(12)]
            out["runs"][regime + "/" + method] = {
                "final_loss_mean": float(np.mean([r["final_loss"] for r in rs])),
                "final_loss_std": float(np.std([r["final_loss"] for r in rs])),
                "p99_update_mean": float(np.mean([r["p99_update"] for r in rs])),
                "divergences": int(sum(r["diverged"] for r in rs)),
                "final_losses": [r["final_loss"] for r in rs]}
    return out

if __name__ == "__main__":
    result = aggregate()
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps({"operator_check": result["operator_check"],
                      "bias_scaling": result["bias_scaling"], "runs": result["runs"]}, indent=2))