#!/usr/bin/env python3
"""Independent numerical reproduction for ORID 69IOkVkTQX.

The implementation follows arXiv:2606.03769v1 directly:
  - LMF: dY=-grad f(X)dt+dL, X=clip(eta(t)Y,-1,1)
  - SDA: Y[t+1]=Y[t]-(grad f(X[t])+noise[t]), X[t]=clip(eta[t]Y[t],-1,1)
  - symmetric Pareto compound-Poisson jumps are centered and have finite p-th
    moments because the tail index is alpha>p; p<2 therefore gives a finite-p,
    potentially infinite-variance regime.

This script is lane-authored and writes only compact JSON/CSV summaries. It
uses scalar streaming updates, one process, and no large trajectory arrays.
"""

from __future__ import annotations

import csv
import json
import math
import os
from pathlib import Path
from typing import Callable

import numpy as np


ROOT = Path(__file__).resolve().parent
OUT = ROOT / "outputs"
OUT.mkdir(parents=True, exist_ok=True)


def huber(x: float, delta: float = 0.25) -> float:
    a = abs(x)
    return 0.5 * x * x / delta if a <= delta else a - 0.5 * delta


def huber_grad(x: float, delta: float = 0.25) -> float:
    a = abs(x)
    if a <= delta:
        return x / delta
    return math.copysign(1.0, x)


def quadratic(x: float) -> float:
    return 0.5 * x * x


def quadratic_grad(x: float) -> float:
    return x


def mirror(z: float, eta: float) -> float:
    return float(np.clip(eta * z, -1.0, 1.0))


def pareto_abs(rng: np.random.Generator, alpha: float, xm: float) -> float:
    u = max(float(rng.random()), 1e-12)
    return xm * u ** (-1.0 / alpha)


def symmetric_pareto(rng: np.random.Generator, alpha: float, xm: float) -> float:
    return (-1.0 if rng.random() < 0.5 else 1.0) * pareto_abs(rng, alpha, xm)


def compound_jump(
    rng: np.random.Generator,
    p: float,
    lam: float = 0.8,
    xm: float = 0.02,
    bounded: float | None = None,
    alpha_offset: float = 0.12,
) -> float:
    alpha = p + alpha_offset
    n = int(rng.poisson(lam))
    total = 0.0
    for _ in range(n):
        mag = pareto_abs(rng, alpha, xm)
        if bounded is not None:
            mag = min(mag, bounded)
        total += (-1.0 if rng.random() < 0.5 else 1.0) * mag
    return total


def pareto_moment(p: float, xm: float = 0.02, alpha_offset: float = 0.12) -> float:
    alpha = p + alpha_offset
    return alpha * (xm**p) / (alpha - p)


def mean_se(values: list[float]) -> tuple[float, float]:
    a = np.asarray(values, dtype=float)
    return float(a.mean()), float(a.std(ddof=1) / math.sqrt(len(a))) if len(a) > 1 else 0.0


def fit_slope(x: list[float], y: list[float]) -> float:
    lx = np.log(np.asarray(x, dtype=float))
    ly = np.log(np.maximum(np.asarray(y, dtype=float), 1e-15))
    return float(np.polyfit(lx, ly, 1)[0])


def fit_seed_slopes(xs: list[float], matrix: list[list[float]]) -> tuple[float, float]:
    slopes = [fit_slope(xs, row) for row in matrix]
    return mean_se(slopes)


def lmf_variable(
    p: float,
    horizon: int,
    seed: int,
    treatment: str = "paper",
    dt: float = 0.5,
    delta: float = 0.25,
    target: float = 0.0,
    reverse_drift: bool = False,
) -> float:
    rng = np.random.default_rng(seed)
    y = 0.0
    avg_x = 0.0
    steps = int(round(horizon / dt))
    for i in range(steps):
        t = (i + 1) * dt
        eta = (1.0 + t) ** (-1.0 / p)
        x = mirror(y, eta)
        avg_x += x * dt
        drift = (1.0 if reverse_drift else -1.0) * huber_grad(x - target, delta) * dt
        if treatment == "deterministic":
            noise = 0.0
        elif treatment == "tame":
            noise = 0.22 * math.sqrt(dt) * float(rng.normal())
        elif treatment == "biased":
            noise = abs(compound_jump(rng, p, lam=0.8, xm=0.02))
        else:
            noise = 0.10 * math.sqrt(dt) * float(rng.normal())
            noise += compound_jump(rng, p, lam=0.8 * dt, xm=0.02)
        y += drift + noise
    return huber(avg_x / horizon - target, delta)


def sda_variable(
    p: float,
    horizon: int,
    seed: int,
    treatment: str = "paper",
    delta: float = 0.25,
    target: float = 0.0,
    reverse_drift: bool = False,
) -> float:
    rng = np.random.default_rng(seed)
    y = 0.0
    avg_x = 0.0
    beta = 0.22
    for t in range(1, horizon + 1):
        eta = beta / (t ** (1.0 / p))
        x = mirror(y, eta)
        avg_x += x
        g = huber_grad(x - target, delta)
        if treatment == "deterministic":
            noise = 0.0
        elif treatment == "biased":
            noise = abs(symmetric_pareto(rng, p + 0.12, 0.03))
        elif p == 2.0:
            noise = 0.18 * float(rng.normal())
        else:
            noise = symmetric_pareto(rng, p + 0.12, 0.03)
        y += g + noise if reverse_drift else -(g + noise)
    return huber(avg_x / horizon - target, delta)


def experiment_claim_1() -> dict:
    ps = [1.25, 1.5, 1.75, 2.0]
    horizons = [256, 512, 1024, 2048, 4096]
    seeds = list(range(101, 109))
    rows = []
    for p in ps:
        matrix = []
        for seed in seeds:
            vals = [lmf_variable(p, T, seed, "paper", delta=0.5, target=0.6) for T in horizons]
            matrix.append(vals)
        slope, slope_se = fit_seed_slopes(horizons, matrix)
        mean_final, se_final = mean_se([row[-1] for row in matrix])
        tamecoef2 = 0.10**2 / 2.0
        heavycoefp = (2.0 ** (2.0 - p)) * 0.8 * pareto_moment(p, 0.02)
        T = horizons[-1]
        exponent = (p - 1.0) / p
        paper_bound = 0.5 / ((1.0 + T) ** exponent) + (p / (p - 1.0)) * tamecoef2 / ((1.0 + T) ** (1.0 / p)) + p * heavycoefp / ((1.0 + T) ** exponent)
        rows.append({
            "p": p,
            "seeds": len(seeds),
            "horizons": horizons,
            "gap_at_max_T": mean_final,
            "gap_se": se_final,
            "measured_slope": slope,
            "slope_se": slope_se,
            "theory_slope": -(p - 1.0) / p,
            "paper_R0_bound_at_max_T": paper_bound,
            "gap_over_paper_bound": mean_final / paper_bound,
            "bound_pass": mean_final <= paper_bound,
        })
    baseline = [lmf_variable(1.5, horizons[-1], s, "deterministic", delta=0.5, target=0.6) for s in seeds]
    destructive = [lmf_variable(1.5, horizons[-1], s, "deterministic", delta=0.5, target=0.6, reverse_drift=True) for s in seeds]
    return {"rows": rows, "baseline": mean_se(baseline), "destructive": mean_se(destructive)}


def lmf_stationary(
    p: float,
    eta: float,
    seed: int,
    treatment: str,
    dt: float = 0.1,
    burn_time: float = 1200.0,
    sample_time: float = 500.0,
) -> float:
    rng = np.random.default_rng(seed)
    y = 0.0
    total_steps = int(round((burn_time + sample_time) / dt))
    burn_steps = int(round(burn_time / dt))
    sumsq = 0.0
    count = 0
    for i in range(total_steps):
        x = mirror(y, eta)
        if i >= burn_steps:
            sumsq += x * x
            count += 1
        noise = 0.0
        if treatment == "tame":
            noise += 0.25 * math.sqrt(dt) * float(rng.normal())
        elif treatment == "heavy":
            noise += compound_jump(rng, p, lam=0.8 * dt, xm=0.02)
        elif treatment == "bounded-control":
            noise += compound_jump(rng, p, lam=0.8 * dt, xm=0.02, bounded=0.12)
        y += -x * dt + noise
    return sumsq / max(count, 1)


def experiment_claim_2() -> dict:
    ps = [1.25, 1.5, 1.75]
    etas = [0.004, 0.008, 0.016, 0.032, 0.064]
    seeds = list(range(201, 209))
    rows = []
    for treatment in ("tame", "heavy", "bounded-control"):
        for p in ps:
            if treatment == "tame" and p != 1.5:
                continue
            matrix = [[lmf_stationary(p, eta, seed, treatment) for eta in etas] for seed in seeds]
            slope, slope_se = fit_seed_slopes(etas, matrix)
            for j, eta in enumerate(etas):
                m, se = mean_se([r[j] for r in matrix])
                rows.append({
                    "treatment": treatment,
                    "p": p,
                    "eta": eta,
                    "seeds": len(seeds),
                    "stationary_mse": m,
                    "mse_se": se,
                    "slope": slope,
                    "slope_se": slope_se,
                    "theory_exponent": 1.0 if treatment != "heavy" else p - 1.0,
                    "paper_radius2_bound": (2.0 * eta * (0.25**2 / 2.0) if treatment == "tame" else 2.0 * (eta ** (p - 1.0)) * (2.0 ** (2.0 - p)) * 0.8 * pareto_moment(p, 0.02)),
                })
                rows[-1]["mse_over_radius_bound"] = rows[-1]["stationary_mse"] / rows[-1]["paper_radius2_bound"]
                rows[-1]["bound_pass"] = rows[-1]["stationary_mse"] <= rows[-1]["paper_radius2_bound"]
    baseline = [lmf_stationary(1.5, 0.016, seed, "none") for seed in seeds]
    return {"rows": rows, "baseline": mean_se(baseline)}


def lmf_hitting(
    p: float,
    delta: float,
    seed: int,
    reverse: bool = False,
    dt: float = 0.2,
    max_time: float = 12000.0,
) -> float | None:
    rng = np.random.default_rng(seed)
    xm = 0.02
    lam = 0.8
    moment_p = lam * pareto_moment(p, xm)
    heavycoef_p = (2.0 ** (2.0 - p)) * moment_p
    eta = (delta * delta / (8.0 * heavycoef_p)) ** (1.0 / (p - 1.0))
    y = 0.90 / eta
    steps = int(round(max_time / dt))
    for i in range(steps + 1):
        t = i * dt
        x = mirror(y, eta)
        if abs(x) <= delta:
            return t
        drift = (1.0 if reverse else -1.0) * x * dt
        y += drift + compound_jump(rng, p, lam=lam * dt, xm=xm)
    return None


def experiment_claim_3() -> dict:
    ps = [1.5, 1.75]
    deltas = [0.50, 0.35, 0.25, 0.18, 0.13]
    seeds = list(range(301, 309))
    rows = []
    for p in ps:
        vals_by_delta = []
        slope_inputs = []
        for delta in deltas:
            vals = [lmf_hitting(p, delta, seed) for seed in seeds]
            finite = [v for v in vals if v is not None]
            mean, se = mean_se(finite) if finite else (float("nan"), float("nan"))
            moment_p = 0.8 * pareto_moment(p, 0.02)
            hc_p = 2.0 ** (2.0 - p) * moment_p
            hc = hc_p ** (1.0 / p)
            bound = 4.0 * 0.5 * ((8.0 ** (1.0 / p) * hc) / (delta * delta)) ** (p / (p - 1.0))
            rows.append({
                "p": p,
                "delta": delta,
                "seeds": len(seeds),
                "mean_hit_time": mean,
                "hit_se": se,
                "censored": len(vals) - len(finite),
                "paper_bound": bound,
                "mean_over_bound": mean / bound if finite else float("nan"),
            })
            if finite:
                slope_inputs.append((delta, mean))
        slope = fit_slope([x for x, _ in slope_inputs], [y for _, y in slope_inputs])
        for row in rows:
            if row["p"] == p:
                row["measured_delta_slope"] = slope
                row["theory_bound_slope"] = -2.0 * p / (p - 1.0)
    baseline = [lmf_hitting(1.5, 0.25, seed, reverse=False) for seed in seeds]
    destructive = [lmf_hitting(1.5, 0.25, seed, reverse=True) for seed in seeds]
    return {"rows": rows, "baseline": baseline, "destructive_censored": sum(v is None for v in destructive)}


def experiment_claim_4() -> dict:
    ps = [1.25, 1.5, 1.75, 2.0]
    horizons = [512, 1024, 2048, 4096, 8192]
    seeds = list(range(401, 409))
    rows = []
    for p in ps:
        sda_matrix = [[sda_variable(p, T, seed, "paper", delta=0.5, target=0.6) for T in horizons] for seed in seeds]
        cont_matrix = [[lmf_variable(p, T, seed, "paper", delta=0.5, target=0.6) for T in horizons] for seed in seeds]
        sda_slope, sda_slope_se = fit_seed_slopes(horizons, sda_matrix)
        cont_slope, cont_slope_se = fit_seed_slopes(horizons, cont_matrix)
        for j, T in enumerate(horizons):
            sda_m, sda_se = mean_se([r[j] for r in sda_matrix])
            cont_m, cont_se = mean_se([r[j] for r in cont_matrix])
            beta = 0.22
            noise_p = pareto_moment(p, 0.03) if p < 2.0 else 0.18**2
            aexp = (p - 1.0) / p
            sda_bound = 0.6**2 / (2.0 * T) + (0.5 + (beta**p) * noise_p) / (beta * (T**aexp))
            lmf_bound = 0.5 / (T**aexp) + p * ((2.0 ** (2.0 - p)) * 0.8 * pareto_moment(p, 0.02)) / (T**aexp)
            rows.append({
                "p": p,
                "T": T,
                "seeds": len(seeds),
                "sda_gap": sda_m,
                "sda_se": sda_se,
                "lmf_gap": cont_m,
                "lmf_se": cont_se,
                "sda_slope": sda_slope,
                "sda_slope_se": sda_slope_se,
                "lmf_slope": cont_slope,
                "lmf_slope_se": cont_slope_se,
                "theory_slope": -(p - 1.0) / p,
                "sda_paper_bound": sda_bound,
                "sda_gap_over_bound": sda_m / sda_bound,
                "sda_bound_pass": sda_m <= sda_bound,
                "lmf_paper_bound": lmf_bound,
                "lmf_gap_over_bound": cont_m / lmf_bound,
                "lmf_bound_pass": cont_m <= lmf_bound,
            })
    baseline = [sda_variable(1.5, horizons[-1], seed, "deterministic", delta=0.5, target=0.6) for seed in seeds]
    destructive = [sda_variable(1.5, horizons[-1], seed, "deterministic", delta=0.5, target=0.6, reverse_drift=True) for seed in seeds]
    return {"rows": rows, "baseline": mean_se(baseline), "destructive": mean_se(destructive)}


def sda_complexity_path(
    p: float,
    eps: float,
    seed: int,
    too_large_eta: bool = False,
    noise_on: bool = True,
    target: float = 0.7,
    reverse_drift: bool = False,
) -> tuple[np.ndarray, float, float]:
    rng = np.random.default_rng(seed)
    sigma_p = pareto_moment(p, 0.03)
    if too_large_eta:
        eta = 0.18 / p
    else:
        eta = min(0.30 / p, 0.15 * (eps / sigma_p) ** (1.0 / (p - 1.0)))
    max_steps = int(max(400, math.ceil(22.0 * math.log(1.0 / eps) / eta)))
    mse = np.empty(max_steps, dtype=float)
    y = 0.0
    for t in range(max_steps):
        x = mirror(y, eta)
        mse[t] = (x - target) * (x - target)
        if noise_on:
            if p == 2.0:
                noise = 0.18 * float(rng.normal())
            else:
                noise = symmetric_pareto(rng, p + 0.12, 0.03)
        else:
            noise = 0.0
        update = (x - target) + noise
        y += update if reverse_drift else -update
    return mse, eta, sigma_p


def crossing(mean_mse: np.ndarray, eps: float) -> int | None:
    window = 40
    if len(mean_mse) < window:
        return None
    rolling = np.convolve(mean_mse, np.ones(window) / window, mode="valid")
    for i, v in enumerate(rolling):
        tail = rolling[i : min(len(rolling), i + 5 * window)]
        if v <= eps and float(np.max(tail)) <= eps * 1.12:
            return i + window
    return None


def experiment_claim_5() -> dict:
    ps = [1.5, 1.75]
    epses = [0.08, 0.05, 0.03, 0.02]
    seeds = list(range(501, 521))
    rows = []
    for p in ps:
        for eps in epses:
            paths = []
            etas = []
            sigma_p = None
            for seed in seeds:
                path, eta, sigma_p = sda_complexity_path(p, eps, seed)
                paths.append(path)
                etas.append(eta)
            n = max(len(x) for x in paths)
            acc = np.zeros(n, dtype=float)
            cnt = np.zeros(n, dtype=float)
            for path in paths:
                acc[: len(path)] += path
                cnt[: len(path)] += 1.0
            mean_mse = acc / np.maximum(cnt, 1.0)
            tstar = crossing(mean_mse, eps)
            tstar_value = float(tstar) if tstar is not None else float("nan")
            ratio = tstar_value / ((eps ** (-1.0 / (p - 1.0))) * math.log(1.0 / eps)) if tstar is not None else float("nan")
            rows.append({
                "p": p,
                "eps": eps,
                "seeds": len(seeds),
                "eta_mean": float(np.mean(etas)),
                "sigma_p": sigma_p,
                "iterations_to_mean_mse_eps": tstar_value,
                "iteration_se": float(np.std([crossing(path, eps) or n for path in paths], ddof=1) / math.sqrt(len(paths))),
                "log_corrected_ratio": ratio,
                "max_mean_mse": float(mean_mse[-1]),
            })
    for p in ps:
        selected = [r for r in rows if r["p"] == p and math.isfinite(r["iterations_to_mean_mse_eps"])]
        slope = fit_slope([r["eps"] for r in selected], [r["iterations_to_mean_mse_eps"] for r in selected])
        for r in selected:
            r["measured_eps_slope"] = slope
            r["theory_power_slope"] = -1.0 / (p - 1.0)
    control_rows = []
    for p in ps:
        for eps in [0.05, 0.02]:
            paths = [sda_complexity_path(p, eps, seed, reverse_drift=True)[0] for seed in seeds]
            final = [float(path[-100:].mean()) for path in paths]
            control_rows.append({
                "p": p,
                "eps": eps,
                "seeds": len(seeds),
                "control_final_mse": mean_se(final)[0],
                "control_final_mse_se": mean_se(final)[1],
                "control_reached": sum(v <= eps for v in final),
            })
    baseline_path = [sda_complexity_path(1.5, 0.03, seed, noise_on=False)[0] for seed in seeds]
    baseline_final = [float(p[-1]) for p in baseline_path]
    return {"rows": rows, "destructive": control_rows, "baseline": mean_se(baseline_final)}


def ito_huber(x: float, a: float) -> float:
    return huber(x, a)


def ito_grad(x: float, a: float) -> float:
    return huber_grad(x, a)


def weak_ito_path(
    seed: int,
    a: float,
    sigma: float = 0.6,
    dt: float = 0.00025,
    T: float = 2.0,
    correction: bool = True,
    kind: str = "huber",
) -> float:
    rng = np.random.default_rng(seed)
    x = 0.0
    rhs = 0.0
    steps = int(round(T / dt))
    for _ in range(steps):
        if kind == "quadratic":
            grad = quadratic_grad(x)
        else:
            grad = ito_grad(x, a)
        dW = sigma * math.sqrt(dt) * float(rng.normal())
        rhs += grad * dW
        if correction:
            rhs += 0.5 * sigma * sigma * dt
        x += dW
        n_large = int(rng.poisson(0.8 * dt))
        for _ in range(n_large):
            j = symmetric_pareto(rng, 1.45, 0.35)
            old = ito_huber(x, a) if kind == "huber" else quadratic(x)
            x += j
            new = ito_huber(x, a) if kind == "huber" else quadratic(x)
            rhs += new - old
    lhs = ito_huber(x, a) if kind == "huber" else quadratic(x)
    return rhs - lhs


def experiment_claim_6() -> dict:
    seeds = list(range(601, 613))
    rows = []
    for a in [1.0, 1.5, 2.0]:
        margins = [weak_ito_path(seed, a, correction=True, kind="huber") for seed in seeds]
        margins_bad = [weak_ito_path(seed, a, correction=False, kind="huber") for seed in seeds]
        m, se = mean_se(margins)
        mb, seb = mean_se(margins_bad)
        rows.append({
            "function": f"Huber(a={a})",
            "seeds": len(seeds),
            "trials": len(seeds),
            "margin_mean": m,
            "margin_se": se,
            "pass_rate": sum(v >= -0.02 for v in margins) / len(margins),
            "destructive_margin_mean": mb,
            "destructive_margin_se": seb,
            "destructive_pass_rate": sum(v >= -0.02 for v in margins_bad) / len(margins_bad),
            "max_jump_note": "unbounded Pareto large jumps, alpha=1.45",
        })
    quad = [weak_ito_path(seed, 0.5, correction=True, kind="quadratic") for seed in seeds]
    return {"rows": rows, "quadratic_baseline": mean_se(quad)}


def write_csv(name: str, rows: list[dict]) -> None:
    if not rows:
        return
    path = OUT / name
    keys = sorted({k for row in rows for k in row})
    with path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=keys)
        w.writeheader()
        w.writerows(rows)


def main() -> None:
    results = {
        "paper": {
            "title": "Bregman meets Lévy: Stochastic mirror descent with heavy-tailed noise in continuous and discrete time",
            "source": "arXiv:2606.03769v1",
            "orid": "69IOkVkTQX",
            "geometry": "X=[-1,1], h(x)=x^2/2, mirror(y)=clip(eta*y,-1,1)",
            "noise": "centered symmetric Pareto / compound-Poisson jumps with alpha=p+0.12; Brownian tame component where stated",
        },
        "claim_1": experiment_claim_1(),
        "claim_2": experiment_claim_2(),
        "claim_3": experiment_claim_3(),
        "claim_4": experiment_claim_4(),
        "claim_5": experiment_claim_5(),
        "claim_6": experiment_claim_6(),
    }
    (OUT / "results.json").write_text(json.dumps(results, indent=2, allow_nan=True), encoding="utf-8")
    for key, value in results.items():
        if key.startswith("claim") and isinstance(value, dict) and isinstance(value.get("rows"), list):
            write_csv(f"{key}.csv", value["rows"])
    print(json.dumps({
        "output": str(OUT / "results.json"),
        "claim_1_rows": len(results["claim_1"]["rows"]),
        "claim_2_rows": len(results["claim_2"]["rows"]),
        "claim_3_rows": len(results["claim_3"]["rows"]),
        "claim_4_rows": len(results["claim_4"]["rows"]),
        "claim_5_rows": len(results["claim_5"]["rows"]),
        "claim_6_rows": len(results["claim_6"]["rows"]),
    }, indent=2))


if __name__ == "__main__":
    main()