import json
import math
import numpy as np

# IMM stale-feedback detector for the known scalar linear process
# z[t+1] = a*z[t] + w[t], y[t] = z[t-delay] + v[t].
# In this verification, the process state history is known to isolate the
# paper's exact likelihood/posterior mechanism from online system-ID error.

def trial(delay_true, noise, a=0.85, q=0.04, D=4, n=1200, seed=0,
          persistence=0.985, threshold=0.01):
    rng = np.random.default_rng(seed)
    z = np.zeros(n + D + 1)
    # Stationary-ish initialization avoids an artificial transient.
    z[0] = rng.normal(0, math.sqrt(q / (1 - a*a)))
    for t in range(n + D):
        z[t + 1] = a * z[t] + rng.normal(0, math.sqrt(q))
    y = np.array([z[D + t - delay_true] + rng.normal(0, math.sqrt(noise))
                  for t in range(n)])

    m = D + 1
    T = np.full((m, m), (1 - persistence) / (m - 1))
    np.fill_diagonal(T, persistence)
    pi = np.ones(m) / m
    posterior = []
    logodds = []
    kl = []
    llr = []
    norm_error = []
    alarm = None
    run = 0
    for t in range(n):
        # IMM mode mixing prior c_i = sum_j T_ji*pi_j.
        c = T.T @ pi
        means = z[D + t - np.arange(m)]
        ll = -0.5 * ((y[t] - means)**2 / noise + math.log(2 * math.pi * noise))
        un = c * np.exp(ll - np.max(ll))
        pi = un / un.sum()
        posterior.append(pi.copy())
        norm_error.append(abs(pi.sum() - 1.0))
        logodds.append(math.log(max(pi[delay_true], 1e-300) /
                              max(pi[0], 1e-300)))
        # Equal-variance Gaussian innovation KL for true vs no-delay mode.
        delta = means[delay_true] - means[0]
        kl.append(delta * delta / (2 * noise))
        # Raw cumulative innovation log-likelihood ratio avoids posterior saturation.
        llr.append(((y[t] - means[0])**2 - (y[t] - means[delay_true])**2) / (2 * noise))
        run = run + 1 if pi[0] <= threshold else 0
        if alarm is None and run >= 3:
            alarm = t + 1

    posterior = np.asarray(posterior)
    logodds = np.asarray(logodds)
    kl = np.asarray(kl)
    llr = np.asarray(llr)
    start = 10
    # Use an early/mid window before numerical posterior saturation.
    stop = min(n, 250)
    slope = float(np.polyfit(np.arange(start, stop), logodds[start:stop], 1)[0])
    cumulative_llr = np.cumsum(llr)
    llr_slope = float(np.polyfit(np.arange(start, stop), cumulative_llr[start:stop], 1)[0])
    mean_kl = float(np.mean(kl[start:stop]))
    dom = next((i + 1 for i, p in enumerate(posterior)
                if p[delay_true] > .5 and p[delay_true] > p[0]), None)
    return dict(slope=slope, llr_slope=llr_slope, mean_kl=mean_kl, ratio=llr_slope / max(mean_kl, 1e-12),
                alarm=alarm, dominant=dom, final_true=float(posterior[-1, delay_true]),
                final_zero=float(posterior[-1, 0]), max_norm_error=max(norm_error))


def aggregate(delay, noise, reps=30, **kw):
    vals = [trial(delay, noise, seed=7000 + k, **kw) for k in range(reps)]
    def med(k):
        x = [v[k] for v in vals if v[k] is not None]
        return None if not x else float(np.median(x))
    return {"delay": delay, "noise": noise, "slope": med("slope"),
            "mean_KL": med("mean_kl"), "cumulative_llr_slope": med("llr_slope"), "slope_over_KL": med("ratio"),
            "alarm_median": med("alarm"), "dominant_median": med("dominant"),
            "final_true_posterior": med("final_true"),
            "final_zero_posterior": med("final_zero"),
            "max_normalization_error": max(v["max_norm_error"] for v in vals)}


def main():
    # Prediction 1: persistent delay separates from no-delay.
    delayed = aggregate(2, .03)
    clean = aggregate(0, .03)
    # Prediction 2: expected log posterior-odds slope is the innovation KL.
    kl_sweep = [aggregate(2, r) for r in (.015, .03, .06, .12)]
    # Prediction 3: lower KL takes longer to cross the fixed alarm threshold.
    detect_sweep = [aggregate(2, r, reps=20) for r in (.015, .03, .06, .12)]
    out = {
      "prediction_checks": {
        "P1_mode_separation": {
          "prediction": "persistent d=2 gives p(d=2)>0.5 while clean d=0 retains p(d=0)",
          "delayed": delayed, "clean": clean},
        "P2_KL_slope": {
          "prediction": "early cumulative innovation log-likelihood slope approximately equals mean innovation KL",
          "noise_sweep": kl_sweep},
        "P3_detection_scaling": {
          "prediction": "alarm time increases as KL decreases (larger observation noise)",
          "noise_sweep": detect_sweep},
        "P4_normalization": {
          "prediction": "posterior sums to one to numerical precision",
          "max_abs_sum_error": max(x["max_normalization_error"] for x in [delayed, clean] + kl_sweep)}
      },
      "settings": {"a": .85, "q": .04, "D": 4, "persistence": .985,
                   "steps": 1200, "threshold": .01, "replicates": 30,
                   "likelihood": "exact Gaussian innovation with known state history"}
    }
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()