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


def f(h, u, lam, q):
    return lam * h + u + q * h * h


def affine_rollout(lam, q, H, gamma):
    # Nominal h*=u*=0, one scalar initial-state coordinate gamma.
    hs = np.zeros(H + 1)
    Rs = np.zeros(H + 1)
    hs[0] = 0.0
    Rs[0] = 1.0
    for k in range(H):
        hs[k + 1] = f(hs[k], 0.0, lam, q)
        # A_k=df/dh at nominal = lam; C_k=1 but U_k=0.
        Rs[k + 1] = lam * Rs[k]
    pred = hs + Rs * gamma
    true = np.zeros(H + 1)
    true[0] = gamma
    for k in range(H):
        true[k + 1] = f(true[k], 0.0, lam, q)
    return hs, Rs, pred, true


def violation(lam, q, H, gamma):
    hs, Rs, pred, true = affine_rollout(lam, q, H, gamma)
    # v_k is one-step nonlinear defect evaluated on the affine family.
    vals = []
    for k in range(H):
        vals.append(abs(f(pred[k], 0.0, lam, q) - pred[k + 1]))
    return max(vals) if vals else 0.0


def max_box_violation(lam, q, H, radius):
    # In this scalar box, the endpoint is worst for this positive quadratic map.
    return max(violation(lam, q, H, -radius), violation(lam, q, H, radius))


def measured_horizon(lam, q, radius, eps, Hmax=80):
    good = 0
    for H in range(1, Hmax + 1):
        if max_box_violation(lam, q, H, radius) <= eps * (1 + 1e-12):
            good = H
        else:
            break
    return good


def predicted_horizon(lam, q, radius, eps):
    # For lambda>1: q*r^2*lambda^(2(H-1)) <= eps.
    if q * radius * radius <= eps:
        return max(1, int(math.floor(1 + math.log(eps/(q*radius*radius)) / (2*math.log(lam)))))
    return 0


def radius_for_horizon(lam, q, H, eps):
    # Exact one-step defect for the affine family at the last step.
    return math.sqrt(eps / (q * lam ** (2 * (H - 1))))


def main():
    np.random.seed(7); random.seed(7)
    lam, q, H, gamma = 1.25, 0.3, 6, 0.04
    # Core math sanity: autodiff-free finite difference Jacobian and recurrence.
    delta = 1e-6
    numeric_A = (f(delta, 0, lam, q) - f(-delta, 0, lam, q)) / (2*delta)
    _, Rs, _, _ = affine_rollout(lam, q, H, gamma)
    jacobian_check = {"analytic_A": lam, "finite_difference_A": numeric_A,
                      "max_R_error": float(np.max(np.abs(Rs - lam**np.arange(H+1))))}

    # Prediction 1: quadratic scaling with perturbation radius.
    radii = np.array([0.01, 0.02, 0.04, 0.08])
    vals = np.array([max_box_violation(lam, q, H, r) for r in radii])
    log_slope = float(np.polyfit(np.log(radii), np.log(vals), 1)[0])

    # Prediction 2: critical horizon follows the exponential boundary.
    eps = 1e-3
    hs = list(range(1, 13))
    measured = [measured_horizon(lam, q, r, eps, 30) for r in [0.01, 0.02, 0.04]]
    predicted = [predicted_horizon(lam, q, r, eps) for r in [0.01, 0.02, 0.04]]

    # Prediction 3: trusted radius decays as lambda^-(H-1).
    horizons = np.array([2, 4, 6, 8])
    rr = np.array([radius_for_horizon(lam, q, int(h), eps) for h in horizons])
    decay_slope = float(np.polyfit(horizons - 1, np.log(rr), 1)[0])
    predicted_decay = -math.log(lam)

    # Parameter sweep over lambda: stable dynamics should not show exponential growth.
    lambda_sweep = []
    for la in [0.8, 1.0, 1.1, 1.25, 1.5]:
        v = max_box_violation(la, q, 10, 0.04)
        lambda_sweep.append({"lambda": la, "H10_violation": v,
                             "critical_H": measured_horizon(la, q, 0.04, eps, 50)})

    # Small monitor comparison: unmonitored affine prediction vs accepted trusted radius.
    # Baseline uses fixed radius 0.04; monitor shrinks to the largest radius satisfying eps.
    baseline_v = max_box_violation(lam, q, 10, 0.04)
    monitored_r = radius_for_horizon(lam, q, 10, eps)
    monitored_v = max_box_violation(lam, q, 10, monitored_r)

    out = {
      "jacobian_check": jacobian_check,
      "prediction_1_quadratic_radius": {"radii": radii.tolist(), "violations": vals.tolist(), "observed_log_slope": log_slope, "predicted_slope": 2.0},
      "prediction_2_horizon_boundary": {"radii": [0.01,0.02,0.04], "epsilon": eps, "measured_H": measured, "predicted_H": predicted, "formula": "q*r^2*lambda^(2(H-1)) <= epsilon"},
      "prediction_3_radius_decay": {"H": horizons.tolist(), "radii": rr.tolist(), "observed_log_decay_per_step": decay_slope, "predicted": predicted_decay},
      "lambda_sweep": lambda_sweep,
      "monitor_comparison": {"baseline_radius": 0.04, "baseline_H10_violation": baseline_v, "monitored_radius": monitored_r, "monitored_H10_violation": monitored_v},
    }
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()