import json
from pathlib import Path
import numpy as np

SEED = 1419
RHO = np.array([1.0, 1.25, 1.5, 2.0, 3.0, 4.0])


def orthonormal_kernel(C):
    # SVD gives an orthonormal basis of ker(C), with C shape (p,n).
    u, s, vh = np.linalg.svd(C, full_matrices=True)
    rank = int(np.sum(s > 1e-10))
    return vh[rank:].T


def projected_sweep(J, C, G0, P):
    U = orthonormal_kernel(C)
    rows = []
    for rho in RHO:
        H = U.T @ (J - rho * G0 @ C) @ U
        S = P @ H + H.T @ P
        rows.append({
            "rho": float(rho),
            "H": H.tolist(),
            "lambda_max_sym": float(np.linalg.eigvalsh((S + S.T) / 2).max()),
        })
    return U, rows


def scheduler(rows, alpha=0.05, P=None):
    # The stated formula uses -2 alpha lambda_min(P).
    pmin = 1.0 if P is None else float(np.linalg.eigvalsh(P).min())
    threshold = -2 * alpha * pmin
    for r in rows:
        if r["lambda_max_sym"] <= threshold:
            return r["rho"]
    return float(rows[-1]["rho"])


def noisy_observer(J, C, G0, rho_policy, steps=3000, dt=0.01,
                   obs_every=5, noise=0.08, seed=0):
    rng = np.random.default_rng(seed)
    n = J.shape[0]
    x = np.array([1.0, 1.0])
    z = np.array([-1.0, 0.0])
    errors = []
    used = []
    for t in range(steps):
        x = x + dt * (J @ x)
        if t % obs_every == 0:
            y = C @ x + noise * rng.normal(size=(C.shape[0],))
            rho = float(rho_policy(z, J, C, G0))
            z = z + dt * (J @ z + rho * G0 @ (y - C @ z))
            used.append(rho)
        else:
            z = z + dt * (J @ z)
        errors.append(float(np.linalg.norm(x - z)))
    tail = np.asarray(errors[-1000:])
    return {"rmse_tail": float(np.sqrt(np.mean(tail ** 2))),
            "mean_rho": float(np.mean(used)), "max_error": float(np.max(errors))}


def main():
    rng = np.random.default_rng(SEED)
    # C observes x1; K is exactly the x2 axis. G0 injects the measurement into x1.
    C = np.array([[1.0, 0.0]])
    G0 = np.array([[1.0], [0.35]])
    P = np.eye(1)
    U = orthonormal_kernel(C)

    # Prediction 1: C U_K=0, hence H(rho) is invariant in rho to machine precision.
    J_stable = np.array([[-0.5, 0.0], [0.0, -0.20]])
    _, stable_rows = projected_sweep(J_stable, C, G0, P)
    Hs = np.asarray([r["H"] for r in stable_rows]).reshape(len(RHO), -1)
    invariant_error = float(np.max(np.abs(Hs - Hs[0])))

    # Prediction 2: the projected boundary is gamma=0.2 (for P=I, alpha=0),
    # independent of gain; the scheduler therefore cannot move it.
    gammas = np.linspace(-0.5, 0.5, 101)
    boundary_records = []
    for gamma in gammas:
        J = np.diag([-0.5, gamma])
        _, rows = projected_sweep(J, C, G0, P)
        boundary_records.append((gamma, rows[0]["lambda_max_sym"]))
    vals = np.array([v for _, v in boundary_records])
    crossing = gammas[np.where(np.sign(vals[:-1]) != np.sign(vals[1:]))[0][0]]

    # Prediction 3: for projected-stable dynamics, scheduler always chooses rho=1,
    # because every candidate has the same certified margin.
    _, stable_rows2 = projected_sweep(J_stable, C, G0, P)
    selected = scheduler(stable_rows2, alpha=0.05, P=P)
    all_selected = []
    for gamma in [-0.5, -0.2, -0.05, 0.05, 0.5]:
        J = np.diag([-0.5, gamma])
        _, rows = projected_sweep(J, C, G0, P)
        all_selected.append({"gamma": gamma, "selected_rho": scheduler(rows, alpha=0.0, P=P)})

    # Comparison: fixed low/high and stated scheduler on the same stable toy system.
    def fixed(rho):
        return lambda z, J, C_, G: rho
    def scheduled(z, J, C_, G):
        _, rows = projected_sweep(J, C_, G, P)
        return scheduler(rows, alpha=0.05, P=P)
    comparison = {
        "fixed_low_rho_1": noisy_observer(J_stable, C, G0, fixed(1.0), seed=SEED),
        "fixed_high_rho_4": noisy_observer(J_stable, C, G0, fixed(4.0), seed=SEED),
        "scheduled": noisy_observer(J_stable, C, G0, scheduled, seed=SEED),
    }

    # Also show that an unstable projected mode cannot be certified at any rho.
    J_unstable = np.diag([-0.5, 0.10])
    _, unstable_rows = projected_sweep(J_unstable, C, G0, P)
    unstable_pick = scheduler(unstable_rows, alpha=0.05, P=P)

    result = {
        "seed": SEED,
        "U_kernel": U.tolist(),
        "prediction_1_invariance_max_abs": invariant_error,
        "prediction_1_expected": 0.0,
        "prediction_2_predicted_boundary_gamma": 0.0,
        "prediction_2_observed_grid_crossing_left_endpoint": float(crossing),
        "prediction_2_boundary_error_grid": float(abs(crossing - 0.0)),
        "prediction_3_stable_selected_rho": selected,
        "prediction_3_expected_rho": 1.0,
        "prediction_3_sweep": all_selected,
        "unstable_projected_mode_scheduler_fallback_rho": unstable_pick,
        "comparison": comparison,
        "interpretation": "C U_K=0 makes the proposed projected gain term identically zero; gain scheduling cannot alter latent projected stability.",
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()