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

SEED = 2015

def lqr(k):
    return (1.0 + k * k) / (2.0 * k)

def lqr_grad(k):
    return 0.5 * (1.0 - 1.0 / (k * k))

def lqr_q(k):
    return abs(lqr_grad(k)) / math.sqrt(max(lqr(k) - 1.0, 1e-30))

def lqr_check():
    # Directly verify the claimed q formula and gradient-flow behavior.
    ks = np.geomspace(1.001, 1000.0, 400)
    formula_err = []
    for k in ks:
        closed = (k + 1.0) / (math.sqrt(2.0) * k ** 1.5)
        formula_err.append(abs(lqr_q(k) - closed) / closed)
    trajectories = []
    for k0 in (1.05, 1.2, 2.0, 10.0):
        eta, k = 1e-4, k0
        gaps, qs, times = [], [], []
        for step in range(30000):
            if step % 100 == 0:
                gaps.append(lqr(k) - 1.0); qs.append(lqr_q(k)); times.append(step * eta)
            k = max(1.0000001, k - eta * lqr_grad(k))
        # local late-trajectory log-gap slope versus the contemporaneous q^2
        y = np.log(np.maximum(gaps, 1e-300)); x = np.asarray(times)
        slope = float(np.polyfit(x[-100:], y[-100:], 1)[0])
        q2 = float(np.mean(np.asarray(qs[-100:]) ** 2))
        trajectories.append({"k0": k0, "log_gap_rate": slope,
                             "minus_mean_q2": -q2,
                             "relative_error": abs(slope + q2) / q2})
    # One-step Euler check of the discrete prediction log gap ~= -eta*q^2.
    discrete = []
    k = 1.01
    for eta in (1e-4, 2e-4, 5e-4, 1e-3, 2e-3):
        old = lqr(k) - 1.0
        knew = k - eta * lqr_grad(k)
        observed = math.log((lqr(knew) - 1.0) / old)
        predicted = -eta * lqr_q(k) ** 2
        discrete.append({"eta": eta, "observed": observed, "predicted": predicted,
                         "relative_error": abs(observed - predicted) / abs(predicted)})
    far_rate = lqr_grad(500.0) ** 2
    return {"max_q_formula_relative_error": float(max(formula_err)),
            "flow_rate_checks": trajectories,
            "discrete_rate_checks": discrete,
            "far_field_rate_at_k500": float(far_rate),
            "checks_pass": bool(max(formula_err) < 1e-8 and
                                  max(r["relative_error"] for r in trajectories) < .01 and
                                  max(r["relative_error"] for r in discrete) < .02 and
                                  abs(far_rate - .25) < .01)}

def make_data(n=1200):
    rng = np.random.RandomState(SEED)
    x = rng.randn(n, 2).astype(np.float32)
    y = (x[:, 0] * x[:, 1] + .25 * x[:, 0] - .15 * x[:, 1] > 0).astype(np.float32)
    return x, y

def mlp_run(scheduled, seed, steps=500, force_cpu=False):
    import torch
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    requested = "cpu" if force_cpu else ("cuda" if torch.cuda.is_available() else "cpu")
    try:
        dev = torch.device(requested)
        x, y = make_data()
        xt, yt = torch.tensor(x[:900], device=dev), torch.tensor(y[:900], device=dev)
        xv, yv = torch.tensor(x[900:], device=dev), torch.tensor(y[900:], device=dev)
        model = torch.nn.Sequential(torch.nn.Linear(2, 32), torch.nn.Tanh(),
                                    torch.nn.Linear(32, 1)).to(dev)
        opt = torch.optim.SGD(model.parameters(), lr=.08)
        lossfn = torch.nn.BCEWithLogitsLoss()
        fhat = None; qhist = []; gaphist = []; consecutive = 0; switched = False
        train, val, lrs, qs = [], [], [], []
        clip_count = 0
        for step in range(steps):
            opt.zero_grad(set_to_none=True)
            loss = lossfn(model(xt).squeeze(-1), yt); loss.backward()
            g = math.sqrt(sum(float((p.grad.detach() ** 2).sum().cpu())
                              for p in model.parameters() if p.grad is not None))
            lv = float(loss.detach().cpu())
            if fhat is None: fhat = lv
            else: fhat = min(fhat, .98 * fhat + .02 * lv)
            gap = max(lv - fhat, 1e-8); q = g / math.sqrt(gap)
            qhist.append(q); gaphist.append(gap); qs.append(q)
            if scheduled and len(qhist) >= 20:
                gaps = np.asarray(gaphist[-100:]); qvals = np.asarray(qhist[-100:])
                candidate = gap <= np.percentile(gaps, 35) and q >= max(np.percentile(qvals, 10) / 2, 1e-5)
                consecutive = consecutive + 1 if candidate else 0
                if consecutive >= 5 and not switched:
                    for group in opt.param_groups: group["lr"] *= 1.5
                    switched = True; switch_step = step
            # Identical clipping in both arms makes the comparison isolate scheduling.
            before = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            clip_count += int(float(before) > 1.0)
            opt.step()
            with torch.no_grad():
                train.append(lv); val.append(float(lossfn(model(xv).squeeze(-1), yv).cpu()))
            lrs.append(opt.param_groups[0]["lr"])
        tail = np.asarray(val[-100:])
        return {"final_train": train[-1], "final_val": val[-1], "best_val": min(val),
                "val_last100_std": float(tail.std()), "switched": switched,
                "switch_step": int(switch_step) if switched else None,
                "lr_final": opt.param_groups[0]["lr"], "clip_steps": clip_count,
                "device": str(dev), "mean_q_last100": float(np.mean(qs[-100:]))}
    except Exception:
        if requested == "cuda":
            try: torch.cuda.empty_cache()
            except Exception: pass
            return mlp_run(scheduled, seed, steps, force_cpu=True)
        raise

def main():
    out = {"math_check": lqr_check(), "mlp": {}}
    for name, scheduled in (("fixed_sgd", False), ("semiglobal_scheduler", True)):
        runs = [mlp_run(scheduled, seed) for seed in (SEED, SEED + 1, SEED + 2)]
        for r in runs: r.pop("device", None)
        for key in ("final_val", "best_val", "val_last100_std", "lr_final", "mean_q_last100"):
            vals = [r[key] for r in runs]
            out["mlp"].setdefault(name, {})[key + "_mean"] = float(np.mean(vals))
            out["mlp"][name][key + "_std"] = float(np.std(vals))
        out["mlp"][name]["runs"] = runs
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()