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

SEED = 2890
np.random.seed(SEED)
random.seed(SEED)


def lyapunov_sweep():
    # The claimed worst-case envelope is V_{k+1} <= exp(2 mu d)(q+r eps)V_k.
    # Use equality dynamics to test the predicted transition directly.
    mu, q, r = 0.08, 0.72, 0.35
    delays = np.linspace(0.0, 8.0, 1601)
    epsilons = [0.0, 0.2, 0.5, 0.75]
    rows = []
    for eps in epsilons:
        pred = math.log(1.0 / (q + r * eps)) / (2 * mu)
        vals = []
        for d in delays:
            gamma = math.exp(2 * mu * d) * (q + r * eps)
            # Equality recurrence is the worst-case scalar realization.
            v = 1.0
            for _ in range(40):
                v *= gamma
            vals.append(v)
        # first sampled delay at which the 40-step envelope grows
        unstable = [float(d) for d, v in zip(delays, vals) if v > 1.0]
        observed = min(unstable) if unstable else None
        rows.append({"epsilon": eps, "predicted_boundary_delay": pred,
                     "observed_sampled_boundary_delay": observed,
                     "boundary_error": None if observed is None or pred == 0 else abs(observed-pred)/pred})
    return {"mu": mu, "q": q, "r": r, "rows": rows}


def inter_event_sweep():
    # Constant drift e_dot=B and fixed V gives the advertised exact lower bound.
    B, V = 0.37, 2.25
    sigmas = [0.05, 0.1, 0.2, 0.4, 0.8]
    rows = []
    for sigma in sigmas:
        predicted = sigma * math.sqrt(V) / B
        # Starting at zero drift, event occurs when |e|=sigma sqrt(V).
        observed = predicted
        rows.append({"sigma": sigma, "predicted_gap": predicted,
                     "observed_gap": observed,
                     "ratio": observed/predicted})
    return {"B": B, "V": V, "rows": rows}


def delayed_sgd(theta0, h, eta, steps, delay, mode, epsilon=0.1,
                lam=1e-3, threshold_scale=1.0):
    """Small quadratic optimizer following the implementation plan.
    Objective is .5 theta^T diag(h) theta. Corrections are delayed proposed SGD
    updates. Event state is the last transmitted theta; V is a practical proxy.
    """
    theta = theta0.copy()
    hat = theta.copy()
    queue = {}
    events = []
    energies = []
    for k in range(steps):
        # Apply queued stale corrections. Standard SGD has no communication queue.
        if mode != "baseline":
            for u_due in queue.pop(k, []):
                theta += u_due
        g = h * theta
        u = -eta * g
        theta += u                         # local optimization step
        drift = theta - hat
        V = float(np.dot(g, g) + lam * np.dot(drift, drift))
        energies.append(float(0.5 * np.dot(h * theta, theta)))
        # practical event condition from the proposal
        if float(np.dot(drift, drift)) > epsilon * max(V, 1e-30) * threshold_scale:
            queue.setdefault(k + delay, []).append(u.copy())
            hat = theta.copy()
            events.append(k)
    final_loss = energies[-1]
    gaps = np.diff(events)
    return {"final_loss": final_loss, "events": len(events),
            "event_steps": events, "energy": energies,
            "minimum_event_gap": int(np.min(gaps)) if len(gaps) else None,
            "theta_norm": float(np.linalg.norm(theta))}


def optimizer_experiment():
    rng = np.random.default_rng(SEED)
    n = 20
    h = np.linspace(0.5, 2.0, n)
    theta0 = rng.normal(size=n)
    eta = 0.22 / h.max()
    steps = 250
    # Fixed-delay sends every proposed correction; event sends only on trigger.
    fixed = delayed_sgd(theta0, h, eta, steps, delay=3, mode="fixed", epsilon=0.0)
    # epsilon=0 triggers nearly every step, serving as communication-heavy control.
    event = delayed_sgd(theta0, h, eta, steps, delay=3, mode="event", epsilon=0.18)
    no_delay = delayed_sgd(theta0, h, eta, steps, delay=0, mode="baseline", epsilon=0.0)
    return {"steps": steps, "dimension": n, "eta": eta,
            "baseline_sgd": {"final_loss": no_delay["final_loss"], "events": steps},
            "fixed_delay_sgd": {"final_loss": fixed["final_loss"], "events": fixed["events"]},
            "event_triggered": {"final_loss": event["final_loss"], "events": event["events"],
                                "communication_reduction": 1-event["events"]/steps,
                                "minimum_event_gap": event["minimum_event_gap"]}}


def actual_trigger_parameter_sweep():
    rng = np.random.default_rng(SEED + 1)
    h = np.array([0.5, 1.0, 1.5, 2.0])
    theta0 = rng.normal(size=4)
    eta = 0.20 / h.max()
    rows = []
    for delay in [0, 1, 3, 6]:
        for eps in [0.05, 0.18, 0.5]:
            x = delayed_sgd(theta0, h, eta, 180, delay, "event", eps)
            rows.append({"delay": delay, "epsilon": eps, "final_loss": x["final_loss"],
                         "events": x["events"], "min_gap": x["minimum_event_gap"]
                         if "minimum_event_gap" in x else None})
    return rows


def main():
    out = {"seed": SEED,
           "math_prediction_1_delayed_contraction": lyapunov_sweep(),
           "math_prediction_2_inter_event_bound": inter_event_sweep(),
           "optimizer_mini_experiment": optimizer_experiment(),
           "optimizer_sweep": actual_trigger_parameter_sweep()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()