import json
from pathlib import Path
import numpy as np
from scipy.optimize import least_squares

DT = 0.02
H = 250
Q_TRUE = 1.0


def carrier(kind="sin", omega=2.0, phase=0.0):
    t = np.arange(H) * DT
    if kind == "sin":
        return np.sin(omega * t + phase)
    if kind == "binary":
        return np.sign(np.sin(omega * t + phase))
    raise ValueError(kind)


def rollout(a, m, q=Q_TRUE, u=None, z0=None):
    if u is None:
        u = np.zeros(H)
    if z0 is None:
        z0 = np.zeros(2)
    z = np.zeros((H + 1, 2))
    z[0] = z0
    for t in range(H):
        inp = u[t] + a * m[t]
        z[t + 1, 0] = z[t, 0] + DT * (-z[t, 0] + q * z[t, 1] ** 2 + inp)
        z[t + 1, 1] = z[t, 1] + DT * (-0.1 * z[t, 1] + inp)
    return z


def gramians(a, m, q=Q_TRUE, delta=1e-12):
    z = rollout(a, m, q)
    # C=[1,0], and A is the exact Jacobian of one Euler step.
    Phi = np.eye(2)
    Wo = np.zeros((2, 2))
    # Reachability is computed by propagating each input impulse forward.
    Wr = np.zeros((2, 2))
    for t in range(H):
        A = np.array([[1 - DT, 2 * DT * q * z[t, 1]],
                      [0, 1 - 0.1 * DT]])
        B = np.array([DT, DT])
        C = np.array([1.0, 0.0])
        Wo += (Phi.T @ np.outer(C, C) @ Phi) * DT
        Psi = np.eye(2)
        for j in range(t + 1, H):
            Aj = np.array([[1 - DT, 2 * DT * q * z[j, 1]],
                           [0, 1 - 0.1 * DT]])
            Psi = Aj @ Psi
        v = Psi @ B
        Wr += np.outer(v, v) * DT
        Phi = A @ Phi
    ev_o = np.linalg.eigvalsh(Wo)
    ev_r = np.linalg.eigvalsh(Wr)
    return Wo, Wr, ev_o, ev_r


def fit_q(observations, m, a):
    # Fit q from observed z1 only, with known input and known initial state.
    def residual(x):
        pred = rollout(a, m, float(x[0]))[:, 0]
        return (pred - observations) / 0.01
    return float(least_squares(residual, [0.2], bounds=(-3, 3)).x[0])


def main():
    rng = np.random.default_rng(123)
    m = carrier("sin", omega=2.0)
    amplitudes = np.array([0.0, 0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.6])
    rows = []
    for a in amplitudes:
        Wo, Wr, eo, er = gramians(a, m)
        rows.append({"a": float(a), "lambda_min_Wo": float(eo[0]),
                     "lambda_min_Wr": float(er[0]),
                     "logdet_Wo": float(np.linalg.slogdet(Wo + 1e-12*np.eye(2))[1])})

    # Prediction 1: passive hidden observability is zero (up to numerical precision).
    passive = rows[0]["lambda_min_Wo"]
    # Prediction 2: small-amplitude hidden observability is quadratic in a.
    small = np.array([r for r in rows if 0.01 <= r["a"] <= 0.08])
    coef = float(np.polyfit(np.log(small[:, 0]) if False else np.log([r["a"] for r in small]),
                            np.log([r["lambda_min_Wo"] for r in small]), 1)[0])
    ratios = [r["lambda_min_Wo"] / (r["a"] ** 2) for r in small]
    # Prediction 3: a practical threshold is where the hidden eigenvalue exceeds 100x passive.
    threshold = next((r["a"] for r in rows if r["lambda_min_Wo"] > max(100*passive, 1e-8)), None)

    # Secondary task check: estimate q from noisy observed trajectories.
    noise = rng.normal(0, 0.01, H + 1)
    passive_obs = rollout(0.0, m)[:, 0] + noise
    probed_a = 0.16
    probed_obs = rollout(probed_a, m)[:, 0] + rng.normal(0, 0.01, H + 1)
    q_passive = fit_q(passive_obs, m, 0.0)
    q_probed = fit_q(probed_obs, m, probed_a)
    result = {"dt": DT, "horizon": H*DT, "rows": rows,
              "predictions": {
                  "passive_lambda_hidden_expected": 0.0,
                  "passive_lambda_observed": passive,
                  "quadratic_exponent_expected": 2.0,
                  "quadratic_exponent_observed": coef,
                  "quadratic_ratio_mean": float(np.mean(ratios)),
                  "quadratic_ratio_cv": float(np.std(ratios)/np.mean(ratios)),
                  "threshold_rule": "lambda_min(Wo) > max(100*passive, 1e-8)",
                  "threshold_observed_a": threshold
              },
              "fit_q": {"true": Q_TRUE, "passive": q_passive, "probed": q_probed}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()