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

SEED = 2027
rng = np.random.default_rng(SEED)
EPS = 1e-8


def throttle(residual, q, kappa, eps=EPS):
    delta = np.linalg.norm(residual) / (math.sqrt(max(q, 0.0)) + eps)
    a = min(1.0, kappa / (delta + eps))
    return float(a), float(delta)


def math_checks():
    # Exact q for J=X. Prediction 1: a transitions at ||r|| = kappa sqrt(q).
    X = np.array([[1.0, 0.0], [0.0, 2.0], [1.0, -1.0]])
    q = float(np.sum(X * X))
    kappa = 0.35
    threshold = kappa * math.sqrt(q)
    residual_norms = np.array([0.25 * threshold, 0.99 * threshold, 1.01 * threshold, 2.0 * threshold, 5.0 * threshold])
    gains = np.array([throttle(np.array([v, 0.0, 0.0]), q, kappa)[0] for v in residual_norms])
    predicted = np.minimum(1.0, threshold / (residual_norms + EPS))
    transition_error = float(np.max(np.abs(gains - predicted)))
    transition_observed = float(residual_norms[np.argmax(gains < 0.999999)])

    # Prediction 2: in the throttled regime, a(c r) * c is constant (inverse gain scaling).
    base = np.array([3.0, -2.0, 1.0])
    scales = np.array([1.0, 2.0, 4.0, 8.0])
    gains_scale = np.array([throttle(c * base, q, kappa)[0] for c in scales])
    inverse_products = scales * gains_scale
    expected_products = np.full_like(inverse_products, kappa * math.sqrt(q) / np.linalg.norm(base))
    # only assess the scales predicted to be throttled
    scale_error = float(np.max(np.abs(inverse_products - expected_products)))

    # Prediction 3: scalar theta'= (1-eta*a(theta))*theta. Plain GD diverges for eta*lambda>2;
    # throttle makes the large-state update approximately eta*kappa and remains bounded.
    eta = 3.0
    kappa1 = 0.4
    n = 80
    theta_plain, theta_throttle = 1.0, 1.0
    plain_path, throttled_path, scalar_gains = [], [], []
    for _ in range(n):
        plain_path.append(abs(theta_plain))
        throttled_path.append(abs(theta_throttle))
        a, _ = throttle(np.array([theta_throttle]), 1.0, kappa1)
        scalar_gains.append(a)
        theta_plain = theta_plain - eta * theta_plain
        theta_throttle = theta_throttle - eta * a * theta_throttle
    plain_growth = abs(theta_plain)
    throttled_final = abs(theta_throttle)
    # In this setup the predicted GD multiplier is |1-eta|=2 and throttle should not grow exponentially.
    expected_plain_growth = 2.0 ** n
    return {
        "q": q,
        "kappa": kappa,
        "transition_predicted_residual_norm": threshold,
        "transition_observed_first_throttled_norm": transition_observed,
        "transition_max_formula_error": transition_error,
        "residual_scales": scales.tolist(),
        "gains_at_scales": gains_scale.tolist(),
        "scale_times_gain": inverse_products.tolist(),
        "scale_invariance_max_error": scale_error,
        "scalar_predicted_plain_multiplier": 2.0,
        "scalar_plain_final_abs": plain_growth,
        "scalar_plain_expected_abs": expected_plain_growth,
        "scalar_throttled_final_abs": throttled_final,
        "scalar_throttled_max_abs": float(max(throttled_path)),
        "scalar_min_gain": float(min(scalar_gains)),
    }


def make_data(n=512, d=8):
    rg = np.random.default_rng(SEED + 1)
    X = rg.normal(size=(n, d))
    true = rg.normal(size=d)
    y = X @ true + 0.03 * rg.normal(size=n)
    return X, y, true


def run_method(method, X, y, true, lr, steps=240, batch=64, kappa=0.5):
    rg = np.random.default_rng(SEED + 10)
    theta = np.zeros(X.shape[1])
    m = np.zeros_like(theta); v = np.zeros_like(theta)
    errors, losses, gains, deltas = [], [], [], []
    for t in range(steps):
        # Alternating masks hide one of two feature groups for 20 steps.
        mask = np.ones(X.shape[1])
        group = (t // 20) % 2
        mask[group::2] = 0.0
        idx = rg.choice(len(X), size=batch, replace=False)
        Xm = X[idx] * mask
        residual = Xm @ theta - y[idx]
        g = Xm.T @ residual / batch
        q = float(np.sum(Xm * Xm) / batch)
        if method == "throttle":
            a, delta = throttle(residual, q, kappa)
            g = a * g
        else:
            a, delta = 1.0, np.linalg.norm(residual) / (math.sqrt(q) + EPS)
        if method == "adam":
            b1, b2 = 0.9, 0.999
            m = b1*m + (1-b1)*g; v = b2*v + (1-b2)*g*g
            step = lr * (m/(1-b1**(t+1))) / (np.sqrt(v/(1-b2**(t+1))) + 1e-8)
            theta -= step
        else:
            theta -= lr * g
        full_res = X @ theta - y
        errors.append(float(np.linalg.norm(theta - true)))
        losses.append(float(np.mean(full_res**2)))
        gains.append(a); deltas.append(delta)
    return {"final_error": errors[-1], "min_error": min(errors), "max_error": max(errors),
            "final_loss": losses[-1], "max_loss": max(losses), "mean_gain": float(np.mean(gains)),
            "fraction_throttled": float(np.mean(np.array(gains) < 0.999999)),
            "error_trace": errors, "loss_trace": losses, "gain_trace": gains, "delta_trace": deltas}


def mini_experiment():
    X, y, true = make_data()
    # This is intentionally near/aggressively beyond the plain linear SGD stability range.
    results = {}
    for lr in [0.08, 0.16, 0.24]:
        results[str(lr)] = {}
        for method in ["sgd", "adam", "throttle"]:
            results[str(lr)][method] = run_method(method, X, y, true, lr=lr, kappa=0.5)
    return results


def main():
    out = {"seed": SEED, "math_checks": math_checks(), "mini_experiment": mini_experiment()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"math_checks": out["math_checks"], "summary": {
        lr: {m: {k: v for k, v in r.items() if k in ["final_error", "max_error", "final_loss", "max_loss", "mean_gain", "fraction_throttled"]}
             for m, r in d.items()} for lr, d in out["mini_experiment"].items()}}, indent=2))

if __name__ == "__main__":
    main()