import json
import numpy as np


def spectral_radius(a):
    return float(np.max(np.abs(np.linalg.eigvals(a))))


def propagate(A, w, e0, horizon):
    e = np.asarray(e0, dtype=float).copy()
    out = [e.copy()]
    for _ in range(horizon):
        e = np.abs(A) @ e + w
        out.append(e.copy())
    return np.asarray(out)


def fit_log_slope(values, start, stop):
    y = np.asarray(values[start:stop])
    x = np.arange(start, stop)
    return float(np.polyfit(x, np.log(np.maximum(y, 1e-300)), 1)[0])


def main():
    np.set_printoptions(precision=6, suppress=True)
    # Positive transition: the absolute-value bound is exact for aligned errors.
    B = np.array([[0.72, 0.18], [0.10, 0.48]], dtype=float)
    rho0 = spectral_radius(B)
    w = np.array([0.01, 0.02])
    e0 = np.array([0.03, 0.01])
    H = 120
    alpha_star = 1.0 / rho0

    # Prediction 1: the boundary is alpha*rho(B)=1.
    alphas = alpha_star * np.array([0.80, 0.95, 0.99, 1.00, 1.01, 1.05, 1.20])
    sweep = []
    for alpha in alphas:
        A = alpha * B
        traj = propagate(A, w, e0, H)
        total = traj.sum(axis=1)
        r = alpha * rho0
        row = {"alpha": float(alpha), "rho_abs": float(r),
               "final_total": float(total[-1]),
               "growth_ratio_last": float(total[-1] / max(total[-2], 1e-30))}
        if r < 1:
            fixed = np.linalg.solve(np.eye(2) - A, w)
            row["predicted_plateau_total"] = float(fixed.sum())
            row["plateau_relative_error"] = float(np.max(np.abs(traj[-1] - fixed)) /
                                                    np.max(np.abs(fixed)))
        else:
            row["predicted_log_growth"] = float(np.log(r))
            row["observed_log_growth"] = fit_log_slope(total, 70, 120)
        sweep.append(row)

    # Prediction 2: below one, finite tube scales as (1-rho)^-1 near boundary.
    near = []
    for r in [0.50, 0.70, 0.80, 0.90, 0.95, 0.98]:
        A = (r / rho0) * B
        fixed = np.linalg.solve(np.eye(2) - A, w)
        near.append({"rho": r, "plateau_total": float(fixed.sum()),
                     "scaled_plateau": float((1-r) * fixed.sum())})

    # Prediction 3: output margin is linear in ebar for y = [1, -0.5] h.
    G = np.array([1.0, 0.5])
    A = 0.75 / rho0 * B
    tr = propagate(A, w, e0, 40)
    margins = tr @ G
    ratios = margins / np.maximum(np.linalg.norm(tr, axis=1, ord=1), 1e-30)
    # Exact validity check with random signed disturbances bounded by w.
    rng = np.random.default_rng(7)
    actual = np.zeros(2)
    bound = np.zeros(2)
    max_violation = 0.0
    for _ in range(1000):
        actual[:] = e0
        bound[:] = e0
        for j in range(30):
            disturbance = rng.uniform(-w, w)
            actual[:] = A @ actual + disturbance
            bound[:] = np.abs(A) @ bound + w
            max_violation = max(max_violation, float(np.max(np.abs(actual) - bound)))

    # Mini-experiment: horizon-dependent local tube versus a single worst-case tube.
    # The latter is the standard conservative replacement A_j -> elementwise max_j A_j.
    Htv = 40
    scales = 0.72 + 0.18 * np.sin(np.arange(Htv) * 0.55) ** 2
    As = [float(sc) / 0.78 * B for sc in scales]
    local = np.zeros(2)
    local_path = [local.copy()]
    for Aj in As:
        local = np.abs(Aj) @ local + w
        local_path.append(local.copy())
    local_path = np.asarray(local_path)
    M = np.max(np.asarray(As), axis=0)
    shared = propagate(M, w, np.zeros(2), Htv)
    local_total = local_path.sum(axis=1)
    shared_total = shared.sum(axis=1)
    result_extra = {
        "time_varying_local_vs_shared": {
            "rho_local_max": float(max(spectral_radius(Aj) for Aj in As)),
            "rho_shared_worst_case": spectral_radius(M),
            "local_final_tube_sum": float(local_total[-1]),
            "shared_final_tube_sum": float(shared_total[-1]),
            "shared_over_local_final": float(shared_total[-1] / local_total[-1]),
            "mean_shared_over_local": float(np.mean(shared_total[1:] / np.maximum(local_total[1:], 1e-30))),
            "local_tube_sums_first_8": [float(x) for x in local_total[:8]],
            "shared_tube_sums_first_8": [float(x) for x in shared_total[:8]]
        }
    }

    result = {
        **result_extra,
        "base_rho": rho0,
        "predicted_alpha_boundary": alpha_star,
        "sweep": sweep,
        "near_boundary_plateau_scaling": near,
        "output_margin_linear_ratio_min": float(np.min(ratios)),
        "output_margin_linear_ratio_max": float(np.max(ratios)),
        "random_disturbance_max_bound_violation": max_violation,
        "formula": "ebar[j+1]=abs(A)@ebar[j]+w; fixed=(I-A)^(-1)w"
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()