import json, math, os
import numpy as np
import torch

SEED = 2244
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)


def device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def omega(rho, curvature, potential=None):
    # A differentiable toy variational functional.  The quartic term makes
    # negative curvature bounded below, while rho=0 exposes its Hessian exactly.
    if potential is None:
        potential = torch.zeros_like(rho)
    return 0.5 * (curvature * rho.square()).sum() + 0.025 * rho.pow(4).sum() + (potential * rho).sum()


def hvp(rho, curvature, v, potential=None, create_graph=False):
    # Keep the graph when the HVP is used inside a trainable regularizer.
    rho_in = rho.detach().clone().requires_grad_(True)
    g = torch.autograd.grad(omega(rho_in, curvature, potential), rho_in, create_graph=True)[0]
    hv = torch.autograd.grad((g * v).sum(), rho_in, create_graph=create_graph)[0]
    return hv if create_graph else hv.detach()


def direct_hessian(rho, curvature, potential=None):
    rho = rho.detach().clone().requires_grad_(True)
    return torch.autograd.functional.hessian(lambda x: omega(x, curvature, potential), rho).detach()


def hess_regularizer(rho, curvature, nvec=16, eps=1e-3, kappa=2.0):
    # Stochastic Rayleigh estimator, exactly in the form proposed in the idea.
    vals = []
    for _ in range(nvec):
        v = torch.randn_like(rho)
        v = v / v.norm()
        vals.append((v * hvp(rho, curvature, v, create_graph=True)).sum())
    q = torch.stack(vals)
    neg = torch.relu(-q + eps).square().mean()
    high = torch.relu(q - kappa).square().mean()
    return neg + high, q.detach()


def run(device_name):
    dev = torch.device(device_name)
    n = 8
    # 1. Core autodiff claim: HVP equals dense Hessian times v.
    c = torch.linspace(0.15, 1.4, n, device=dev)
    r = torch.randn(n, device=dev) * 0.2
    v = torch.randn(n, device=dev)
    H = direct_hessian(r, c)
    hv_err = (hvp(r, c, v) - H @ v).abs().max().item()
    eigs = torch.linalg.eigvalsh(H).cpu().numpy()

    # 2. Susceptibility sweep. For quadratic Omega, drho/dV=-H^-1 exactly.
    gammas = np.array([0.02, 0.04, 0.08, 0.16, 0.32, 0.64])
    suscept = []
    predicted = []
    for g in gammas:
        # use a fixed unit perturbation so the norm is analytically 1/gamma
        H_g = torch.eye(n, device=dev) * float(g)
        dv = torch.zeros(n, device=dev); dv[0] = 1.0
        dr = torch.linalg.solve(H_g, -dv)
        suscept.append(dr.norm().item()); predicted.append(1.0 / g)
    slope = float(np.polyfit(np.log(gammas), np.log(suscept), 1)[0])
    rel_err = float(np.max(np.abs(np.array(suscept)-predicted) / np.array(predicted)))

    # 3. Zero crossing: explicit Euler response iteration is stable iff gamma*lambda < 2.
    # For the fixed-point minimization rho <- rho - eta(H rho + dV), eta=1.
    # Its amplification factor is |1-gamma|; instability starts at gamma=2.
    cross_g = np.array([0.5, 1.5, 2.0, 2.5])
    growth = []
    for g in cross_g:
        x = 1.0
        for _ in range(30): x = (1.0-g)*x
        growth.append(abs(x))
    # Negative curvature is directly visible as negative Rayleigh quotient.
    neg_c = torch.tensor(-0.2, device=dev)
    neg_rho = torch.zeros(8, device=dev)
    reg_neg, q_neg = hess_regularizer(neg_rho, neg_c, nvec=32, eps=1e-3, kappa=2.0)

    # 4. Mini experiment: fit an observed curvature target that is spuriously negative.
    # Baseline can retain the negative mode; regularized fit rejects it.
    target = -0.20
    def fit(use_reg, target_value=target, steps=300):
        x = torch.tensor(-0.5, device=dev, requires_grad=True)
        opt = torch.optim.Adam([x], lr=0.03)
        rr = torch.zeros(8, device=dev)
        for _ in range(steps):
            opt.zero_grad()
            loss = (x-target_value).square()
            if use_reg:
                reg, _ = hess_regularizer(rr, x, nvec=16, eps=1e-3, kappa=2.0)
                loss = loss + 8.0 * reg
            loss.backward(); opt.step()
        return float(x.detach().cpu()), float((x-target_value)**2)
    base, base_fit = fit(False)
    regular, reg_fit = fit(True)

    # Deliberately near-critical positive mode should not be pushed to kappa;
    # only the negative-curvature branch is active.
    soft_target = 0.02
    soft_base, _ = fit(False, soft_target)
    soft_reg, _ = fit(True, soft_target)

    result = {
        "device": str(dev),
        "hvp_max_abs_error": hv_err,
        "hessian_eigenvalues": eigs.tolist(),
        "susceptibility_sweep": [{"lambda": float(g), "observed": float(s), "predicted_1_over_lambda": float(p)} for g,s,p in zip(gammas,suscept,predicted)],
        "loglog_slope_observed": slope,
        "loglog_slope_predicted": -1.0,
        "max_relative_scaling_error": rel_err,
        "euler_zero_crossing_predicted_gamma": 2.0,
        "euler_growth_sweep": [{"gamma": float(g), "abs_amplification_after_30": float(a)} for g,a in zip(cross_g,growth)],
        "negative_mode_mean_rayleigh": float(q_neg.mean().cpu()),
        "negative_mode_regularizer": float(reg_neg.cpu()),
        "curvature_fit": {"target_negative": target, "baseline_final": base, "regularized_final": regular, "baseline_mse": base_fit, "regularized_mse": reg_fit, "soft_target": soft_target, "soft_baseline_final": soft_base, "soft_regularized_final": soft_reg},
    }
    return result


def main():
    try:
        d = device()
        out = run(d)
    except Exception as e:
        if torch.cuda.is_available():
            out = run("cpu")
            out["cuda_error_fallback"] = repr(e)
        else:
            raise
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()