import json, math, random
import numpy as np


def c_transition(k):
    k = max(float(k), 1.0)
    return math.sqrt(((k + 1.0) + math.sqrt(max(0.0, (k + 1.0)**2 - 4.0))) / 2.0)


def calibrated_t(k, delta):
    k = max(float(k), 1.0)
    delta = float(delta)
    t = math.sqrt(1.0 + math.sqrt(max(0.0, (k - 1.0) * (1.0 / delta - 1.0))))
    return max(t, c_transition(k))


def tail_formula(t, k):
    return (k - 1.0) / ((t*t - 1.0)**2 + k - 1.0) if k > 1 else 0.0


def math_check():
    rows = []
    max_inv_err = 0.0
    min_tail_margin = float('inf')
    for k in [1.01, 1.1, 2, 3, 10, 100]:
        c = c_transition(k)
        poly = c**4 - (k + 1)*c**2 + 1
        for d in [1e-3, .01, .1, .5]:
            raw = math.sqrt(1 + math.sqrt((k-1)*(1/d-1)))
            t = calibrated_t(k, d)
            # Inversion is exact before regime projection; the clipping projection is deliberate.
            if raw >= c:
                inv_err = abs(tail_formula(raw, k) - d)
                max_inv_err = max(max_inv_err, inv_err)
            min_tail_margin = min(min_tail_margin, t-c)
        rows.append({'k': k, 'c': c, 'transition_polynomial': poly})
    return {'max_inverse_error': max_inv_err, 'minimum_projected_tail_margin': min_tail_margin, 'rows': rows}


def sample_kurtosis(x):
    z = x - np.mean(x)
    v = np.mean(z*z)
    return float(np.mean(z**4) / max(v*v, 1e-18))


def empirical_tail_check(seed=17, n=2000000):
    rng = np.random.default_rng(seed)
    # Standardized Student-t has a known kurtosis 3 + 6/(nu-4), nu>4.
    out = []
    for nu in [5, 8, 20]:
        x = rng.standard_t(nu, n)
        x = (x - x.mean()) / x.std()
        k_bound = 3 + 6/(nu-4)
        for d in [.01, .05]:
            t = calibrated_t(k_bound*1.05, d)  # modest safety factor
            observed = float(np.mean(x >= t))
            out.append({'distribution': 'student_t', 'nu': nu, 'kurtosis_bound': k_bound*1.05,
                        'delta': d, 'threshold': t, 'observed_tail': observed,
                        'bound_formula_at_threshold': tail_formula(t, k_bound*1.05)})
    return out


def run_optimizer(seed, mode, steps=300, batch=64, delta=.02):
    rng = np.random.default_rng(seed)
    x = 8.0
    beta = .95
    mu_ema, v_ema, q_ema = 0., 1., 3.
    losses, spikes, tails = [], 0, []
    lr = .08
    for step in range(steps):
        # Rare, centered, heavy-tailed gradient contamination.
        noise = rng.normal(0, .35, batch)
        rare = rng.random(batch) < .025
        noise[rare] += rng.choice([-1, 1], rare.sum()) * 12.0
        g = x + noise
        if mode == 'calibrated':
            m = float(g.mean())
            z = g - m
            vv = float(np.mean(z*z))
            qq = float(np.mean(z**4))
            mu_ema = beta*mu_ema + (1-beta)*m
            v_ema = beta*v_ema + (1-beta)*vv
            q_ema = beta*q_ema + (1-beta)*qq
            k = min(1000., max(1., 1.5*q_ema/(v_ema*v_ema + 1e-12)))
            t = calibrated_t(k, delta)
            tau = t*math.sqrt(v_ema + 1e-12)
            clipped = m + np.clip(g-m, -tau, tau)
            used = float(clipped.mean())
            tails.append(float(np.mean(np.abs(g-m) > tau)))
            spikes += int(np.any(np.abs(g-m) > tau))
        elif mode == 'fixed':
            tau = 2.0
            used = float(np.clip(g, -tau, tau).mean())
            tails.append(float(np.mean(np.abs(g) > tau)))
            spikes += int(np.any(np.abs(g) > tau))
        else:
            used = float(g.mean())
            spikes += int(np.any(np.abs(g) > 8.0))
        x -= lr * used
        losses.append(.5*x*x)
    return {'final_loss': losses[-1], 'mean_last50_loss': float(np.mean(losses[-50:])),
            'max_loss': max(losses), 'spike_steps': spikes,
            'mean_tail_above_threshold': float(np.mean(tails)) if tails else None,
            'final_abs_x': abs(x)}


def main():
    results = {'math_check': math_check(), 'empirical_tail_check': empirical_tail_check()}
    allruns = {}
    for mode in ['none', 'fixed', 'calibrated']:
        vals = [run_optimizer(s, mode) for s in range(10, 20)]
        allruns[mode] = {k: (None if vals[0][k] is None else float(np.mean([v[k] for v in vals]))) for k in vals[0]}
        allruns[mode]['replicate_std_final_loss'] = float(np.std([v['final_loss'] for v in vals]))
    results['optimizer'] = allruns
    with open('results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print(json.dumps(results, indent=2))

if __name__ == '__main__':
    main()