"""Small noisy snapshot comparison: unfiltered versus forward-compatible EDMD."""
import json
import numpy as np
from forward_intersection import make_K, principal_filter


def one_trial(seed, T=80, noise=0.04, tau=0.20, horizon=10):
    rng = np.random.default_rng(seed)
    Q0 = np.eye(5)[:, :4]
    K = make_K(.60)
    # Snapshot inputs lie in the dictionary; outputs are noisy ambient observations.
    X = rng.normal(size=(4, T))
    Y = K @ (Q0 @ X) + noise * rng.normal(size=(5, T))
    Khat = Y @ np.linalg.pinv(Q0 @ X)
    Q1, _ = principal_filter(Q0, Khat @ Q0, tau)
    A0 = Q0.T @ Khat @ Q0
    Af = Q1.T @ Khat @ Q1

    # Held-out initial conditions contain only represented stable modes.
    Z = rng.normal(size=(3, 30))
    truth_diag = np.array([.9, .7, .8])
    err0, err1 = [], []
    for h in range(1, horizon + 1):
        target = truth_diag ** h @ Z
        pred0 = (A0[:3, :3] ** 0) @ Z  # overwritten below; explicit matrix power follows
        pred0 = np.linalg.matrix_power(A0[:3, :3], h) @ Z
        # The selected Q1 is normally the first three coordinate directions; use
        # its ambient prediction and compare only to the true supported trajectory.
        pred1 = Q1 @ np.linalg.matrix_power(Af, h) @ (Q1.T @ (Q0[:,:3] @ Z))
        err0.append(np.mean((pred0 - target) ** 2))
        err1.append(np.mean((pred1[:3] - target) ** 2))
    return {
        "dim": int(Q1.shape[1]),
        "spectral_radius_baseline": float(max(abs(np.linalg.eigvals(A0)))),
        "spectral_radius_idea": float(max(abs(np.linalg.eigvals(Af)))),
        "rollout_mse_10_baseline": float(err0[-1]),
        "rollout_mse_10_idea": float(err1[-1]),
    }


def run():
    results = [one_trial(s) for s in range(10)]
    keys = ["spectral_radius_baseline", "spectral_radius_idea",
            "rollout_mse_10_baseline", "rollout_mse_10_idea"]
    means = {k: float(np.mean([r[k] for r in results])) for k in keys}
    return {"settings": {"trials": 10, "T": 80, "noise_std": .04,
                          "tau": .20, "horizon": 10},
            "mean": means,
            "dimensions": [r["dim"] for r in results], "trials": results}


if __name__ == "__main__":
    print(json.dumps(run(), indent=2))