import json, math
from pathlib import Path
import numpy as np
from scipy.linalg import expm
from scipy.stats import spearmanr

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

# Stable damped oscillator: x(t)=exp(A t)x(0), y=c^T x.
A = np.array([[-0.08, -2.0], [2.0, -0.08]], dtype=float)
c = np.array([1.0, 0.0])

def observability(delays):
    return np.stack([c @ expm(-A * float(t)) for t in delays])

def metrics(O):
    s = np.linalg.svd(O, compute_uv=False)
    return float(s[-1]), float(s[0]), float(s[0]/s[-1])

def reconstruct(O, x, sigma, reps=3000):
    # Gaussian measurement noise; average squared state error.
    pinv = np.linalg.pinv(O)
    noise = rng.normal(0.0, sigma, size=(reps, O.shape[0]))
    errs = (noise @ pinv.T)
    rmse = float(np.sqrt(np.mean(np.sum(errs*errs, axis=1))))
    return rmse

def main():
    out = {"seed": SEED, "A": A.tolist(), "c": c.tolist(), "predictions": {}, "designs": [], "noise_sweep": [], "spearman": {}}

    # Prediction 1: isotropic LS RMSE is sigma*sqrt(trace((O^T O)^-1)),
    # hence linear in noise level (and bounded by sqrt(n)*sigma/smin).
    delays = np.array([0.0, 0.30, 0.63, 0.95, 1.30])
    O = observability(delays)
    smin, smax, cond = metrics(O)
    trace_factor = math.sqrt(float(np.trace(np.linalg.inv(O.T @ O))))
    noise_levels = [0.002, 0.005, 0.01, 0.02, 0.05, 0.1]
    vals = []
    for sig in noise_levels:
        actual = reconstruct(O, None, sig)
        predicted = sig * trace_factor
        bound = math.sqrt(2.0) * sig / smin
        vals.append({"sigma": sig, "observed_rmse": actual, "predicted_rmse": predicted, "worst_case_bound": bound})
    slopes = np.polyfit(np.log(noise_levels), np.log([v['observed_rmse'] for v in vals]), 1)[0]
    out["predictions"]["noise_linear_scaling"] = {"predicted_slope": 1.0, "observed_loglog_slope": float(slopes), "design_smin": smin, "design_condition": cond}
    out["noise_sweep"] = vals

    # Prediction 2: for fixed white noise, MSE^0.5 follows the exact inverse Gram factor,
    # and should be strongly monotone with ill-conditioning across delay sets.
    designs = {
        "uniform_short": np.linspace(0, 0.35, 5),
        "uniform_medium": np.linspace(0, 1.30, 5),
        "uniform_long": np.linspace(0, 2.60, 5),
        "irregular_good": np.array([0.0, 0.22, 0.61, 1.07, 1.56]),
        "irregular_bad_clustered": np.array([0.0, 0.02, 0.04, 0.07, 0.11]),
        "irregular_multiscale": np.array([0.0, 0.08, 0.40, 1.30, 2.35]),
    }
    sigma = 0.03
    factors, observed = [], []
    for name, d in designs.items():
        oo = observability(d)
        sm, sx, kk = metrics(oo)
        factor = math.sqrt(float(np.trace(np.linalg.inv(oo.T @ oo))))
        obs = reconstruct(oo, None, sigma)
        factors.append(factor); observed.append(obs)
        out["designs"].append({"name": name, "delays": d.tolist(), "smin": sm, "smax": sx, "condition": kk, "predicted_rmse": sigma*factor, "observed_rmse": obs})
    rho, p = spearmanr(factors, observed)
    out["spearman"] = {"rho_predicted_factor_vs_observed_rmse": float(rho), "p_value": float(p)}
    out["predictions"]["inverse_singular_value_noise_amplification"] = {"expected": "positive monotonic relation", "observed_spearman": float(rho), "threshold": 0.7}

    # Prediction 3: adding a non-redundant delay can improve smin; sweep a fifth delay.
    # Compare a clustered baseline against candidate irregular placement.
    base = np.array([0.0, 0.25, 0.55, 0.85])
    candidates = np.linspace(0.0, 3.0, 301)
    smins = np.array([metrics(observability(np.r_[base, t]))[0] for t in candidates])
    best_i = int(np.argmax(smins))
    out["predictions"]["delay_design_sweep"] = {
        "predicted": "non-redundant delay increases smallest singular value",
        "base_delays": base.tolist(), "base_smin": metrics(observability(base))[0],
        "best_added_delay": float(candidates[best_i]), "best_smin": float(smins[best_i]),
        "improvement_ratio": float(smins[best_i]/metrics(observability(base))[0]),
        "candidate_grid_step": 0.01
    }

    # Mini baseline comparison: actual delays are Poisson-irregular. The baseline
    # decodes with a fixed uniform-delay observation matrix, while the idea uses
    # the observed timestamps to build O for each history.
    nominal = np.array([0.0, 0.30, 0.60, 0.90, 1.20])
    baseline_O = observability(nominal)
    poisson_rows = []
    for mean_gap in [0.08, 0.15, 0.30, 0.50]:
        base_errors, conditioned_errors = [], []
        for _ in range(1200):
            # Past gaps, newest observation at delay zero.
            gaps = rng.exponential(mean_gap, size=4)
            actual_delays = np.r_[0.0, np.cumsum(gaps)]
            actual_O = observability(actual_delays)
            noise = rng.normal(0.0, 0.03, size=5)
            x_true = rng.normal(size=2)
            y = actual_O @ x_true + noise
            base_errors.append(np.linalg.norm(np.linalg.pinv(baseline_O) @ y - x_true))
            conditioned_errors.append(np.linalg.norm(np.linalg.pinv(actual_O) @ y - x_true))
        poisson_rows.append({
            "mean_gap": mean_gap,
            "baseline_fixed_uniform_error": float(np.mean(base_errors)),
            "conditioned_actual_delay_error": float(np.mean(conditioned_errors)),
            "relative_improvement": float(1.0 - np.mean(conditioned_errors)/np.mean(base_errors))
        })
    out["poisson_baseline_comparison"] = poisson_rows

    # A direct bound check over random noise vectors: ||pinv(O)e|| <= ||e||/smin.
    bound_ratios = []
    for _ in range(10000):
        e = rng.normal(size=O.shape[0])
        bound_ratios.append(np.linalg.norm(np.linalg.pinv(O) @ e) / (np.linalg.norm(e)/smin))
    out["bound_check"] = {"max_ratio": float(max(bound_ratios)), "mean_ratio": float(np.mean(bound_ratios)), "claim": "ratio <= 1"}

    # Pass criteria: mechanism checks, not merely a baseline win.
    out["worked_checks"] = {
        "noise_slope_within_0.10": bool(abs(slopes-1.0) <= 0.10),
        "spearman_above_0.7": bool(rho >= 0.7),
        "bound_holds": bool(max(bound_ratios) <= 1.0 + 1e-10),
        "delay_improves_smin": bool(smins[best_i] > metrics(observability(base))[0] * 1.05),
    }
    out["worked"] = all(out["worked_checks"].values())
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()