import json
import numpy as np
from scipy.special import ndtr
from reachable_risk import reachable_risk, fixed_false_warning_recall

SEED = 1449
rng = np.random.default_rng(SEED)


def mechanism_sweeps():
    # Prediction 1: q=Phi(-m), hence q(0)=.5 and q falls monotonically with m.
    margins = np.array([-3., -2., -1., 0., 1., 2., 3.])
    q = ndtr(-margins)
    monotone = bool(np.all(np.diff(q) < 0))
    q_at_zero = float(q[margins == 0][0])

    # Prediction 2: scalar covariance follows s_h=lambda^2 s_(h-1)+Q.
    # For lambda<1 it saturates at Q/(1-lambda^2); lambda>1 grows exponentially.
    Q, s0, H = .04, .01, 20
    growth = {}
    for lam in [.5, .9, 1.0, 1.1, 1.4]:
        s = s0
        vals = []
        for _ in range(H):
            s = lam * lam * s + Q
            vals.append(s)
        growth[str(lam)] = {
            "s_h": vals,
            "ratio_last_first": float(vals[-1] / vals[0]),
            "stable_limit_if_applicable": (Q / (1-lam*lam)) if abs(lam) < 1 else None,
        }
    stable_sat_error = abs(growth["0.5"]["s_h"][-1] - Q/(1-.5**2))
    unstable_ratio = growth["1.4"]["ratio_last_first"]

    # Prediction 3: for a fixed positive per-step q, union risk is 1-(1-q)^H.
    q0 = .1
    horizons = np.arange(1, 11)
    union = 1 - (1-q0) ** horizons
    exact_h5 = float(union[4])

    return {
        "margin_sweep": {"margins": margins.tolist(), "q": q.tolist(),
                          "monotone_decreasing": monotone, "q_at_m_zero": q_at_zero,
                          "predicted_q_at_m_zero": .5},
        "covariance_sweep": growth,
        "covariance_predictions": {
            "stable_lambda_0.5_limit": Q/(1-.5**2),
            "observed_lambda_0.5_final": growth["0.5"]["s_h"][-1],
            "stable_absolute_error": float(stable_sat_error),
            "unstable_lambda_1.4_ratio_last_first": unstable_ratio,
            "predicted_unstable_ratio_gt_1": True,
        },
        "horizon_sweep": {"q_per_step": q0, "horizons": horizons.tolist(),
                          "union_risk": union.tolist(), "risk_at_h5": exact_h5,
                          "predicted_h5": 1-(1-q0)**5},
    }


def rollout_trial(x0, H, noise_sd, threshold):
    # state=(position, velocity), x_{t+1}=(position+velocity, velocity)+epsilon
    x = np.array(x0, dtype=float)
    states = []
    for _ in range(H):
        x = np.array([x[0] + x[1], x[1]]) + rng.normal(0, noise_sd, 2)
        states.append(x.copy())
    return np.asarray(states), bool(np.any(np.asarray(states)[:, 0] >= threshold))


def mini_experiment(n=2500, H=8):
    # Every episode has a fresh initial state; warning is generated before rollout.
    J = np.array([[1., 1.], [0., 1.]])
    Q = np.diag([.025**2, .012**2])
    init_cov = np.diag([.03**2, .02**2])
    halfspaces = [(np.array([1., 0.]), 1.0)]
    noise_sd = np.array([.025, .012])
    risk_scores, mean_scores, labels = [], [], []
    for _ in range(n):
        x0 = np.array([rng.uniform(.05, .85), rng.uniform(.01, .13)])
        controls = [None] * H
        def dyn(mu, u):
            return J @ mu, J
        risk = reachable_risk(x0, init_cov, dyn, controls, halfspaces, Q)
        # Fair deterministic baseline: mean-only endpoint warning probability
        # with a fixed noise scale, without reachable-set covariance growth.
        mean_future = np.array([x0[0] + H*x0[1], x0[1]])
        mean_scores.append(float(ndtr((mean_future[0] - 1.0) / 0.05)))
        risk_scores.append(float(risk))
        _, violated = rollout_trial(x0, H, noise_sd, 1.0)
        labels.append(violated)
    labels = np.asarray(labels)
    risk_scores, mean_scores = np.asarray(risk_scores), np.asarray(mean_scores)
    rr, rt = fixed_false_warning_recall(risk_scores, labels, .05)
    mr, mt = fixed_false_warning_recall(mean_scores, labels, .05)
    # 10 equal-count bins give a simple empirical expected calibration error.
    def ece(scores):
        order = np.argsort(scores); chunks = np.array_split(order, 10); out = 0.
        for ix in chunks:
            out += len(ix)/len(scores) * abs(np.mean(labels[ix])-np.mean(scores[ix]))
        return float(out)
    return {"n": n, "horizon": H, "violation_rate": float(np.mean(labels)),
            "reachable_risk_recall_at_5pct_fpr": rr,
            "mean_score_recall_at_5pct_fpr": mr,
            "reachable_threshold": rt, "mean_threshold": mt,
            "reachable_risk_ece": ece(risk_scores), "mean_score_ece": ece(mean_scores),
            "baseline_definition": "fixed-sigma endpoint Gaussian warning probability (sigma=0.05)"}


if __name__ == "__main__":
    result = {"seed": SEED, "mechanism": mechanism_sweeps(), "mini_experiment": mini_experiment()}
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))