import json
import numpy as np
from pathlib import Path


def transport_queue_check(seed=0):
    rng = np.random.default_rng(seed)
    d = 7
    u = rng.normal(size=40)
    # Discrete upwind transport: q_{k+1}[s]=q_k[s-1], q_k[0]=u_k.
    q = np.zeros(d)
    observed = []
    for k, uk in enumerate(u):
        q[1:] = q[:-1].copy()
        q[0] = uk
        observed.append(q[-1])
    expected = [0.0] * (d - 1) + list(u[:-(d - 1)])
    err = float(np.max(np.abs(np.asarray(observed) - np.asarray(expected))))
    return {"delay": d, "max_abs_error": err}


def closed_form_predictor(A, B, x, queued):
    d = len(queued)
    out = np.linalg.matrix_power(A, d) @ x
    for j, uj in enumerate(queued):
        out += np.linalg.matrix_power(A, d - 1 - j) @ B @ uj
    return out


def predictor(A, B, x, queued):
    """Compute A^d x + sum A^(d-1-j) B queued[j]."""
    p = np.asarray(x, dtype=float).copy()
    for uj in queued:  # queued is chronological: oldest input first
        p = A @ p + B @ np.asarray(uj, dtype=float)
    return p


def simulate(A, B, K, d, steps=250, seed=123, compensated=True, noise_std=0.0):
    rng = np.random.default_rng(seed)
    n, m = B.shape
    x = np.array([1.0, -0.7], dtype=float)
    # pipeline entries are controls which will be applied over the next d steps
    queue = [np.zeros(m) for _ in range(d)]
    norms = []
    controls = []
    for k in range(steps):
        queued = [q.copy() for q in queue]
        applied = queue.pop(0)
        if compensated:
            target = predictor(A, B, x, queued)
        else:
            target = x
        u = K @ target
        queue.append(np.asarray(u).copy())
        x = A @ x + B @ applied + noise_std * rng.normal(size=n)
        norms.append(float(np.linalg.norm(x)))
        controls.append(float(np.linalg.norm(u)))
    return {
        "rms_last_50": float(np.sqrt(np.mean(np.square(norms[-50:])))),
        "max_norm": float(np.max(norms)),
        "final_norm": float(norms[-1]),
        "norms": norms,
        "controls": controls,
    }


def lyapunov_margin(F):
    # P from the discrete Lyapunov series; this is a constructive numerical check.
    n = F.shape[0]
    P = np.zeros((n, n))
    term = np.eye(n)
    for _ in range(10000):
        P += term
        term = F.T @ term @ F
        if np.linalg.norm(term) < 1e-12:
            break
    eig = np.linalg.eigvalsh(P)
    residual = F.T @ P @ F - P
    # For Q=I, residual should be -I up to truncation.
    return {"rho": float(np.max(np.abs(np.linalg.eigvals(F)))),
            "P_min_eig": float(np.min(eig)),
            "lyapunov_max_eig": float(np.max(np.linalg.eigvalsh(residual))),
            "series_terms": _}


def main():
    # A mildly unstable plant and a gain with a stable no-delay target.
    A = np.array([[1.08, 0.12], [0.0, 1.03]])
    B = np.array([[0.0], [1.0]])
    K = np.array([[-0.72, -0.48]])
    F = A + B @ K
    rng = np.random.default_rng(9)
    Aq = np.array([[0.91, 0.13], [-0.04, 0.86]])
    Bq = np.array([[0.2], [0.7]])
    xq = rng.normal(size=2)
    uq = [rng.normal(size=1) for _ in range(5)]
    algebraic_error = float(np.max(np.abs(predictor(Aq, Bq, xq, uq) - closed_form_predictor(Aq, Bq, xq, uq))))
    out = {"transport_check": transport_queue_check(), "predictor_formula_check": {"max_abs_error": algebraic_error}, "target": lyapunov_margin(F), "runs": []}
    # Equal setup, fixed seeds; baseline is stale-state feedback.
    for d in [1, 2, 4, 8, 16, 32]:
        base = simulate(A, B, K, d, compensated=False)
        pred = simulate(A, B, K, d, compensated=True)
        out["runs"].append({"d": d, "baseline": {k: v for k, v in base.items() if k not in ("norms", "controls")},
                            "predictor": {k: v for k, v in pred.items() if k not in ("norms", "controls")}})
    # Noise sensitivity is reported at a representative long delay.
    out["noise"] = []
    for sigma in [0.0, 0.005, 0.02]:
        b = simulate(A, B, K, 16, compensated=False, noise_std=sigma)
        p = simulate(A, B, K, 16, compensated=True, noise_std=sigma)
        out["noise"].append({"sigma": sigma, "baseline_rms_last_50": b["rms_last_50"],
                             "predictor_rms_last_50": p["rms_last_50"]})
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()