import json
from pathlib import Path
import numpy as np


def ar1_stats(a, sigma=1.0, n=30000, burn=3000, seed=0):
    rng = np.random.default_rng(seed)
    x = np.zeros(n + burn)
    for t in range(n + burn - 1):
        x[t + 1] = a * x[t] + sigma * rng.normal()
    y = x[burn:]
    return float(np.var(y)), float(np.corrcoef(y[:-1], y[1:])[0, 1])


class CSDMonitor:
    def __init__(self, window=80, ac_threshold=0.94, gamma=0.9):
        self.window = window
        self.ac_threshold = ac_threshold
        self.gamma = gamma
        self.values = []
        self.previous_variance = None
        self.last = {"a": np.nan, "variance": np.nan, "trigger": False}

    def update(self, value):
        self.values.append(float(value))
        if len(self.values) < self.window:
            return 1.0, False
        z = np.asarray(self.values[-self.window:])
        z = z - z.mean()
        variance = float(np.mean(z * z))
        denom = float(np.dot(z[:-1], z[:-1]))
        a_hat = float(np.dot(z[:-1], z[1:]) / denom) if denom > 1e-12 else 0.0
        rising = self.previous_variance is not None and variance > self.previous_variance
        trigger = bool(a_hat > self.ac_threshold and rising)
        self.previous_variance = variance
        self.last = {"a": a_hat, "variance": variance, "trigger": trigger}
        multiplier = float(np.exp(-self.gamma * max(a_hat - self.ac_threshold, 0.0))) if trigger else 1.0
        return multiplier, trigger


def variance_and_ac_checks():
    # Prediction 1: stationary variance is sigma^2/(1-a^2).
    # Prediction 2: lag-1 correlation equals a.
    rows = []
    for i, a in enumerate([0.20, 0.50, 0.70, 0.85, 0.93]):
        var, rho = ar1_stats(a, seed=100 + i)
        pred_var = 1.0 / (1.0 - a * a)
        rows.append({"a": a, "predicted_variance": pred_var, "observed_variance": var,
                     "predicted_rho": a, "observed_rho": rho,
                     "variance_relative_error": abs(var - pred_var) / pred_var,
                     "rho_abs_error": abs(rho - a)})
    return rows


def boundary_check():
    # Prediction 3: deterministic AR(1) is stable below a=1 and grows above a=1.
    # Estimate the boundary by classifying a from the final/initial amplitude ratio.
    candidates = np.linspace(0.94, 1.06, 25)
    classified = []
    for a in candidates:
        x = 1.0
        for _ in range(80):
            x *= a
        classified.append(abs(x) > 1.0)
    transitions = [candidates[i] for i in range(1, len(candidates))
                   if classified[i] != classified[i - 1]]
    observed = float(transitions[0]) if transitions else float("nan")
    return {"predicted_boundary": 1.0, "observed_grid_boundary": observed,
            "grid_step": float(candidates[1] - candidates[0])}


def intervention_experiment(seed=2024, trials=120):
    # A slowly worsening plant: a ramps from .88 to 1.035. The monitor acts by
    # reducing effective gain. Both policies see exactly the same noise per trial.
    rng = np.random.default_rng(seed)
    horizon = 320
    window = 40
    ac_threshold = 0.90
    gamma = 1.8
    base_rms, safe_rms, base_danger, safe_danger, leads = [], [], [], [], []
    for _ in range(trials):
        noise = rng.normal(size=horizon) * 0.25
        a_schedule = np.linspace(0.90, 1.12, horizon)
        xb = xs = 0.0
        monitor = CSDMonitor(window, ac_threshold, gamma)
        xs_hist = []
        triggered_at = None
        for t in range(horizon - 1):
            xb = a_schedule[t] * xb + noise[t]
            xs_hist.append(xs)
            mult, trigger = monitor.update(xs)
            # Damping is applied to the unstable plant coefficient, clipped to
            # preserve a positive recovery coefficient.
            effective_a = max(0.0, a_schedule[t] * mult)
            xs = effective_a * xs + noise[t]
            if trigger and triggered_at is None:
                triggered_at = t
        base_rms.append(float(np.sqrt(np.mean(np.asarray([0.0] + xs_hist) ** 2))))
        safe_rms.append(float(np.sqrt(np.mean(np.asarray(xs_hist) ** 2))))
        base_danger.append(float(np.max(np.abs(np.asarray([0.0] + xs_hist))) > 4.0))
        safe_danger.append(float(np.max(np.abs(np.asarray(xs_hist))) > 4.0))
        if triggered_at is not None:
            # independently measured crossing is first index with a >= 1
            crossing = int(np.argmax(a_schedule >= 1.0))
            leads.append(crossing - triggered_at)
    return {
        "trials": trials, "baseline_rms": float(np.mean(base_rms)),
        "monitor_rms": float(np.mean(safe_rms)),
        "baseline_danger_rate": float(np.mean(base_danger)),
        "monitor_danger_rate": float(np.mean(safe_danger)),
        "trigger_rate": len(leads) / trials,
        "mean_warning_lead_steps": float(np.mean(leads)) if leads else None,
        "median_warning_lead_steps": float(np.median(leads)) if leads else None,
        "settings": {"window": window, "a_c": ac_threshold, "gamma": gamma}
    }


def main():
    checks = variance_and_ac_checks()
    boundary = boundary_check()
    intervention = intervention_experiment()
    max_var_error = max(r["variance_relative_error"] for r in checks)
    max_rho_error = max(r["rho_abs_error"] for r in checks)
    report = {
        "variance_ac_checks": checks,
        "boundary_check": boundary,
        "intervention": intervention,
        "pass_criteria": {
            "max_variance_relative_error_lt_0.10": max_var_error < 0.10,
            "max_rho_abs_error_lt_0.02": max_rho_error < 0.02,
            "boundary_within_grid_step": abs(boundary["observed_grid_boundary"] - 1.0) <= boundary["grid_step"]
        },
        "interpretation": "The first two checks test stationary AR(1) formulas; the intervention is a small illustrative safety simulation, not a trained controller benchmark."
    }
    Path("csd_results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()