import json
from pathlib import Path
import numpy as np

SEED = 1339
DT = 0.002
T = 2.0
TS = np.arange(0.0, T + DT / 2, DT)
ALPHA = 0.10
EPS = 0.02


def smooth_positive(y, width=31):
    width = int(width) | 1
    pad = width // 2
    yp = np.pad(y, (pad, pad), mode="edge")
    return np.convolve(yp, np.ones(width) / width, mode="valid") + 1e-8


def demonstrations(rng, n=160):
    x0 = rng.uniform(-1.0, 1.0, n)
    k = rng.normal(1.0, 0.07, n)
    x = x0[:, None] * np.exp(-k[:, None] * TS[None, :])
    x += rng.normal(0.0, 0.006, x.shape)
    return x


def infer_funnel(xdemo):
    raw = EPS + np.quantile(np.abs(xdemo), 1.0 - ALPHA, axis=0)
    return smooth_positive(raw)


def rollout(x0, ubar, k_request, rho, aware):
    x = float(x0)
    xs, us, gains = [], [], []
    for j in range(len(TS)):
        r = float(rho[j])
        k = min(k_request, ubar / r) if aware else k_request
        u = float(np.clip(-k * x, -ubar, ubar))
        xs.append(x)
        us.append(u)
        gains.append(k)
        x += DT * u
    return np.asarray(xs), np.asarray(us), np.asarray(gains)


def fit_decay(x):
    sel = (TS >= 0.15) & (TS <= 0.8) & (np.abs(x) > 1e-5)
    return -float(np.polyfit(TS[sel], np.log(np.abs(x[sel])), 1)[0])


def main():
    rng = np.random.default_rng(SEED)
    demos = demonstrations(rng)
    rho = infer_funnel(demos)
    # Prediction 1: actuator-aware requested gain is clipped at k_max=ubar/rho.
    # Use the initial funnel radius, where the scalar bound is sharpest.
    rho0 = float(rho[0])
    ubar = 0.35
    predicted_kmax = ubar / rho0
    gain_rows = []
    for requested in [0.5 * predicted_kmax, 0.9 * predicted_kmax,
                      1.1 * predicted_kmax, 2.0 * predicted_kmax]:
        _, _, gains = rollout(1.0, ubar, requested, rho, True)
        gain_rows.append({"requested": requested,
                          "observed_initial_gain": float(gains[0]),
                          "ratio_to_bound": float(gains[0] / predicted_kmax)})

    # Prediction 2: below the bound, normalized error decays as exp(-k t),
    # while above it actuator clipping causes a slower initial decay.
    decay_rows = []
    for requested in [0.5 * predicted_kmax, 0.9 * predicted_kmax,
                      1.1 * predicted_kmax, 2.0 * predicted_kmax]:
        for aware in [False, True]:
            x, u, gains = rollout(1.0, ubar, requested, rho, aware)
            decay_rows.append({"requested": requested, "aware": aware,
                               "predicted_decay": requested,
                               "observed_decay": fit_decay(x),
                               "max_abs_u": float(np.max(np.abs(u))),
                               "max_normalized_error": float(np.max(np.abs(x) / rho)),
                               "violation_fraction": float(np.mean(np.abs(x) > rho))})

    # Prediction 3: increasing actuator authority increases the feasible gain
    # linearly, k_max=ubar/rho0, and the aware policy realizes that scaling.
    scaling_rows = []
    requested = 10.0
    for limit in [0.15, 0.25, 0.35, 0.50, 0.70]:
        _, _, gains = rollout(1.0, limit, requested, rho, True)
        predicted = limit / rho0
        scaling_rows.append({"ubar": limit, "predicted_kmax": predicted,
                             "observed_initial_gain": float(gains[0]),
                             "relative_error": float(abs(gains[0] - predicted) / predicted)})

    # Held-out initial states: compare fixed-gain saturated baseline to funnel-aware law.
    heldout = np.linspace(-1.0, 1.0, 41)
    k_baseline = 2.0 * predicted_kmax
    comparisons = []
    for aware in [False, True]:
        vals = []
        for x0 in heldout:
            x, u, _ = rollout(x0, ubar, k_baseline, rho, aware)
            vals.append([abs(x[-1]), np.max(np.abs(x) / rho), np.mean(np.abs(x) > rho), np.mean(np.abs(u))])
        a = np.asarray(vals)
        comparisons.append({"controller": "funnel_aware" if aware else "fixed_gain_saturated",
                            "final_abs_error_mean": float(a[:, 0].mean()),
                            "max_normalized_error_mean": float(a[:, 1].mean()),
                            "violation_fraction_mean": float(a[:, 2].mean()),
                            "mean_abs_control": float(a[:, 3].mean())})

    # Cheap numerical derivative check of the normalized-error identity.
    x, _, _ = rollout(0.8, ubar, 0.5 * predicted_kmax, rho, True)
    z = x / rho
    zd_num = np.gradient(z, DT)
    edot = np.gradient(x, DT)
    rhodot = np.gradient(rho, DT)
    zd_formula = edot / rho - x * rhodot / (rho ** 2)
    identity_rel_error = float(np.max(np.abs(zd_num - zd_formula)) / (np.max(np.abs(zd_num)) + 1e-9))

    results = {
        "seed": SEED,
        "n_demonstrations": len(demos),
        "rho_initial": rho0,
        "rho_final": float(rho[-1]),
        "predictions": {
            "actuator_boundary": {"predicted_kmax_ubar_over_rho": predicted_kmax, "sweep": gain_rows},
            "decay_and_saturation": decay_rows,
            "linear_authority_scaling": scaling_rows,
            "normalized_error_identity_relative_error": identity_rel_error,
        },
        "comparison": comparisons,
    }
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()