import json
from pathlib import Path
import numpy as np

SEED = 1729


def sliced_w2_sq(a, b, n_proj=128, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    d = a.shape[1]
    q = rng.normal(size=(n_proj, d))
    q /= np.linalg.norm(q, axis=1, keepdims=True)
    pa = np.sort(a @ q.T, axis=0)
    pb = np.sort(b @ q.T, axis=0)
    return float(d * np.mean((pa - pb) ** 2))


def exact_translation_sweep():
    """Verify W2^2 <= D dt Sigma for a distribution translated at velocity u.

    For p_t(x)=p_0(x-u t), v=u, so W2^2=||u||^2 dt^2 and
    D dt Sigma = D dt * (||u||^2/D)dt = ||u||^2 dt^2 exactly.
    This avoids confusing paired particle displacement with distribution W2.
    """
    rng = np.random.default_rng(SEED)
    rows = []
    for d in (1, 2, 8):
        for D in (0.05, 0.2, 0.8):
            for dt in (0.02, 0.08):
                for speed in (0.4, 1.2):
                    n = 1024
                    x0 = rng.normal(size=(n, d))
                    u = np.full(d, speed / np.sqrt(d))
                    x1 = x0 + dt * u
                    w2 = sliced_w2_sq(x0, x1, 128, rng)
                    sigma = float(np.dot(u, u) / D)
                    rhs = D * dt * sigma * dt
                    predicted = float(np.dot(u, u) * dt * dt)
                    rows.append({"d": d, "D": D, "dt": dt, "speed": speed,
                                 "w2": w2, "rhs": rhs, "predicted": predicted,
                                 "ratio": w2 / rhs,
                                 "w2_over_speed2_dt2": w2 / (speed * speed * dt * dt)})
    ratios = np.array([r["ratio"] for r in rows])
    scaling = np.array([r["w2_over_speed2_dt2"] for r in rows])
    return rows, {"median_ratio": float(np.median(ratios)),
                  "max_abs_ratio_error": float(np.max(np.abs(ratios - 1))),
                  "median_speed_dt_scaling": float(np.median(scaling)),
                  "speed_dt_scaling_range": [float(np.min(scaling)), float(np.max(scaling))]}


def controller_update(eta, ratio, eta_min, eta_max, alpha=0.35, target=0.72):
    # Safety-oriented interpretation of the prose: r>target decreases eta;
    # r<target permits a cautious increase, with multiplicative clipping.
    factor = np.clip((target / max(ratio, 1e-8)) ** alpha, 0.5, 1.25)
    return float(np.clip(eta * factor, eta_min, eta_max))


def quadratic_run(eta0, controlled, seed, steps=240, K=32, d=4):
    rng = np.random.default_rng(seed)
    H = np.diag(np.linspace(0.5, 4.0, d))
    x = rng.normal(0, 2, size=(K, d))
    eta = eta0
    T = 0.002
    D = T
    eta_min, eta_max = eta0 * 0.01, eta0 * 2.0
    losses, etas, rs, disp = [], [], [], []
    old = x.copy()
    for t in range(steps):
        grad = x @ H.T
        x = x - eta * grad + np.sqrt(2 * eta * T) * rng.normal(size=x.shape)
        losses.append(float(np.mean(0.5 * np.sum((x @ H) * x, axis=1))))
        etas.append(eta)
        if (t + 1) % 8 == 0:
            interval = 8 * eta
            delta = x - old
            sigma = float(np.mean((delta / interval) ** 2) / D)
            rhs = D * interval * sigma * interval
            w2 = sliced_w2_sq(old, x, 64, rng)
            ratio = w2 / max(rhs, 1e-12)
            rs.append(ratio)
            disp.append(w2)
            if controlled:
                eta = controller_update(eta, ratio, eta_min, eta_max)
            old = x.copy()
    return {"final_loss": losses[-1], "min_loss": min(losses),
            "max_loss": max(losses), "final_eta": eta,
            "mean_ratio": float(np.mean(rs)), "max_ratio": float(np.max(rs)),
            "losses": losses, "etas": etas, "ratios": rs, "displacements": disp}


def main():
    rows, toy_summary = exact_translation_sweep()
    sweep = []
    for eta in (0.10, 0.30, 0.45, 0.55, 0.70):
        sweep.append({"eta0": eta,
                      "baseline": quadratic_run(eta, False, SEED),
                      "controller": quadratic_run(eta, True, SEED)})
    out = {
        "seed": SEED,
        "toy": {"rows": rows, "summary": toy_summary},
        "quadratic": sweep,
        "predictions": {
            "P1_bound": "For translation, W2^2/(D dt Sigma)=1 independent of D.",
            "P2_scaling": "W2^2 scales as speed^2*dt^2; the dimension-corrected sliced estimator is invariant to D and dimension for fixed total speed.",
            "P3_optimizer_boundary": "For H eigenvalue 4, noiseless GD boundary is eta=2/4=0.5; high-eta baseline should diverge and controller should reduce eta."
        }
    }
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"toy_summary": toy_summary,
      "quadratic_summary": [{"eta0": r["eta0"],
        "base_final": r["baseline"]["final_loss"],
        "ctrl_final": r["controller"]["final_loss"],
        "base_max_ratio": r["baseline"]["max_ratio"],
        "ctrl_max_ratio": r["controller"]["max_ratio"],
        "ctrl_final_eta": r["controller"]["final_eta"]} for r in sweep]}, indent=2))

if __name__ == "__main__":
    main()