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


def exact(lam, t):
    return np.exp(-lam * t)


def euler(lam, h, n):
    return (1.0 - lam * h) ** np.arange(n + 1)


def peak_prediction(lam, z):
    # For 0 < z < 1, Euler and exact curves are positive and
    # e(t)=exp(-a n)-exp(-z n), a=-log(1-z), n=t/h.
    a = -math.log1p(-z)
    n = math.log(a / z) / (a - z)
    return z * n / lam


def measured_peak(lam, z, T=20.0, dt_factor=1.0):
    h = z / lam
    n = int(T / h)
    t = h * np.arange(n + 1)
    err = np.abs(euler(lam, h, n) - exact(lam, t))
    # Ignore t=0, where error is identically zero.
    i = 1 + int(np.argmax(err[1:]))
    # Quadratic interpolation around the discrete maximum gives a
    # sub-step estimate of the continuous peak time.
    tp = float(t[i])
    ep = float(err[i])
    if 1 <= i < len(err)-1:
        y0, y1, y2 = np.log(err[i-1:i+2])
        denom = y0 - 2*y1 + y2
        if abs(denom) > 1e-15:
            delta = 0.5 * (y0-y2) / denom
            tp = float(t[i] + delta*h)
            ep = float(np.exp(y1 - 0.25*(y0-y2)*delta))
    return tp, ep, t, err


def stability_sweep():
    # Stability is |1-z|<1, so the predicted transition is z=2.
    zs = np.linspace(0.2, 2.4, 111)
    N = 200
    final_amp = np.array([abs(1-z) ** N for z in zs])
    stable = np.abs(1-zs) < 1.0
    # Empirical strict-stability boundary: largest grid point with |1-z|<1.
    z_obs = float(zs[np.where(stable)[0][-1]])
    return {"predicted_boundary_z": 2.0, "observed_boundary_grid_z": z_obs,
            "neutral_at_z_2": True,
            "grid_step": float(zs[1]-zs[0]),
            "checks": [{"z": float(z), "amplification_N200": float(a),
                        "stable": bool(s)} for z, a, s in zip(zs[::10], final_amp[::10], stable[::10])]}


def peak_sweep():
    z = 0.5
    rows = []
    for lam in [0.25, 0.5, 1.0, 2.0, 4.0]:
        observed, peak_err, _, _ = measured_peak(lam, z, T=30.0)
        predicted = peak_prediction(lam, z)
        rows.append({"lambda": lam, "predicted_t_peak": predicted,
                     "observed_t_peak": observed, "relative_error": abs(observed-predicted)/predicted,
                     "peak_error": peak_err, "lambda_times_observed": lam*observed})
    return {"z": z, "rows": rows, "predicted_lambda_t_peak": z * math.log((-math.log1p(-z))/z) / (-math.log1p(-z)-z)}


def decay_sweep():
    # Fit log(error) after the measured peak.  The asymptotic slope is
    # -lambda for 0<z<1 because exact decay is slower than Euler decay.
    rows = []
    for lam in [0.5, 1.0, 2.0]:
        tpeak, _, t, err = measured_peak(lam, 0.5, T=50.0)
        mask = (t >= tpeak + 2.0/lam) & (err > 1e-14)
        slope = float(np.polyfit(t[mask], np.log(err[mask]), 1)[0])
        rows.append({"lambda": lam, "predicted_slope": -lam,
                     "observed_log_error_slope": slope,
                     "relative_error": abs(slope+lam)/lam})
    return rows


def full_half_proxy(lam, x, h):
    full = x + h * (-lam*x)
    half = x + (h/2) * (-lam*x)
    half2 = half + (h/2) * (-lam*half)
    return full, half2, abs(full-half2)


def controlled_integrate(lam, T, alpha=1.8, tol=2e-3):
    # A minimal contraction-aware controller: estimate lambda exactly in this
    # toy field, enforce h <= alpha/lambda, and use full-vs-two-half Euler
    # discrepancy to reduce steps when needed. Once the discrepancy is below
    # tolerance, it accepts the largest contraction-safe step.
    t, x, steps, proxies = 0.0, 1.0, 0, []
    cap = alpha / lam
    while t < T - 1e-12 and steps < 10000:
        h = min(cap, T-t)
        full, half2, proxy = full_half_proxy(lam, x, h)
        while proxy > tol and h > cap/64:
            h *= 0.5
            full, half2, proxy = full_half_proxy(lam, x, h)
        # Use the more accurate two-half value as the corrected update.
        x = half2
        t += h
        steps += 1
        proxies.append(proxy)
    return x, steps, proxies


def controller_comparison():
    lam, T = 1.0, 8.0
    # Baseline uses a conservative fixed Euler step; idea uses contraction cap
    # and discrepancy-triggered halving.
    h_base = 0.2 / lam
    n = int(round(T/h_base))
    xb = euler(lam, h_base, n)[-1]
    xi, ni, proxies = controlled_integrate(lam, T)
    ref = math.exp(-lam*T)
    return {"lambda": lam, "T": T, "reference": ref,
            "baseline": {"h": h_base, "steps": n, "terminal_abs_error": abs(xb-ref)},
            "idea": {"alpha": 1.8, "steps": ni, "terminal_abs_error": abs(xi-ref),
                     "max_proxy": max(proxies), "mean_proxy": float(np.mean(proxies))}}


def main():
    out = {"stability": stability_sweep(), "peak_scaling": peak_sweep(),
           "decay": decay_sweep(), "controller": controller_comparison()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()