import json
import numpy as np
from scipy.optimize import minimize

SEED = 1436
rng = np.random.default_rng(SEED)


def sector_matrix(alpha=0.0, beta=1.0):
    return np.array([[-alpha * beta, 0.5 * (alpha + beta)],
                     [0.5 * (alpha + beta), -1.0]])


def solve_certificate(A, B1, B2, C1, D11, D12, C2, D21, D22,
                      alpha=-1.0, beta=1.0, eps=1e-7):
    """Fast scalar certificate search: exact LMI eigenvalues on a log grid."""
    X=np.array([[float(np.asarray(A).ravel()[0]),float(B1),float(B2)]])
    V=np.array([[float(C1),float(D11),float(D12)]])
    E=np.array([[float(C2),float(D21),float(D22)]])
    W=np.array([[0.,1.,0.]])
    Q=np.vstack([V,W]).T @ sector_matrix(alpha,beta) @ np.vstack([V,W])
    best=None
    # Homogeneity permits P=1 in the scalar state dimension; scan multiplier and gain.
    for P in np.logspace(-3,3,9):
     for lam in np.r_[0., np.logspace(-2,2,8)]:
      for g2 in np.r_[np.logspace(-2,4,15)]:
        L=X.T@(P*X)-np.diag([P,0.,0.])+E.T@E-g2*np.diag([0.,0.,1.])+lam*Q
        ev=float(np.linalg.eigvalsh(L).max())
        if ev < -eps and (best is None or g2 < best[0]): best=(g2,lam,ev,P)
    if best is None: return False,np.inf,np.nan,np.nan,np.inf
    return True,float(np.sqrt(best[0])),float(best[3]),float(best[1]),float(best[2])


def collect_and_identify(a, b, n=5000, noise=0.0):
    """Exciting trajectory and least-squares reconstruction of [A B1 B2], [C1 D11 D12]."""
    u = rng.normal(0.0, 0.35, n)
    x = np.zeros(n + 1); w = np.zeros(n); v = np.zeros(n)
    for k in range(n):
        v[k] = x[k]
        w[k] = np.tanh(v[k])
        x[k + 1] = a*x[k] + b*w[k] + u[k] + noise*rng.normal()
    # exact hidden state is available in this controlled verification; identify linear lifted model
    Z = np.column_stack([x[:-1], w, u])
    Xhat = np.linalg.lstsq(Z, x[1:], rcond=None)[0]
    Vhat = np.linalg.lstsq(Z, v, rcond=None)[0]
    return Xhat, Vhat, (x, u, v, w)


def nonlinear_metrics(a, b, n=250, delta=1e-5):
    """Local perturbation decay and finite-horizon impulse gain."""
    x1 = delta; x2 = -delta
    diffs = []
    for _ in range(n):
        diffs.append(abs(x1-x2))
        x1 = a*x1 + b*np.tanh(x1)
        x2 = a*x2 + b*np.tanh(x2)
    ratio = (diffs[-1] / diffs[0]) ** (1.0 / max(n-1, 1))
    # disturbance impulse, zero initial state; output is x
    x = 0.0; ss = 0.0
    for k in range(n):
        d = 1.0 if k == 0 else 0.0
        x = a*x + b*np.tanh(x) + d
        ss += x*x
    return ratio, np.sqrt(ss)


def main():
    # b is fixed and a is swept. Tanh has sector [-1,1]; robust sector boundary is |a|+b=1.
    b = 0.45
    a_values = np.array([-1.45,-1.1,-0.8,-0.5,-0.2,0.0,0.2,0.4,0.5,0.6,0.75])
    rows = []
    for a in a_values:
        cert = solve_certificate(a, b, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0)
        identX, identV, _ = collect_and_identify(a, b, n=500)
        data_cert = solve_certificate(identX[0], identX[1], identX[2],
                                      identV[0], identV[1], identV[2],
                                      1.0, 0.0, 0.0)
        rho, impulse = nonlinear_metrics(a, b)
        rows.append({"a": float(a), "local_rho_pred": abs(a+b),
                     "measured_rho": float(rho), "impulse_l2": float(impulse),
                     "cert": bool(cert[0]), "gamma": float(cert[1]),
                     "data_cert": bool(data_cert[0]), "data_gamma": float(data_cert[1]),
                     "lmi_maxeig": float(cert[4])})

    # A finer boundary search reports the first certified stable point and compares it
    # with the analytical local and sector worst-case boundaries.
    fine = np.linspace(-1.2, 0.7, 10)
    fine_cert = [solve_certificate(a, b, 1.0, 1.0, 0, 0, 1, 0, 0)[0] for a in fine]
    cert_stable = [a for a, ok in zip(fine, fine_cert) if ok]
    observed_boundary = max([a for a in cert_stable if a >= 0.0], default=np.nan)
    # For tanh sector [-1,1], the positive-a robust boundary is a+b=1.
    predicted_boundary = 1.0 - b
    # Fit gain scaling in stable region against exact linearized Hinf prediction 1/(1-rho).
    stable_rows = [r for r in rows if r["local_rho_pred"] < 0.92 and r["cert"]]
    for r in stable_rows:
        r["linear_gain_pred"] = 1.0 / (1.0 - r["local_rho_pred"])
    out = {"seed": SEED, "b": b, "rows": rows,
           "prediction_boundary_a": predicted_boundary,
           "observed_certificate_boundary_a": float(observed_boundary),
           "boundary_error": float(observed_boundary-predicted_boundary),
           "positive_boundary_definition": "largest nonnegative certified a on the tested grid",
           "stable_gain_rows": stable_rows,
           "notes": [
             "Prediction 1: local perturbations contract for |a+b|<1 and grow for |a+b|>1.",
             "Prediction 2: tanh sector certificate boundary is a+b=1 on the positive-a side.",
             "Prediction 3: disturbance gain grows as 1/(1-(a+b)) near the boundary."
           ]}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps({"predicted_boundary": predicted_boundary,
                      "observed_certificate_boundary": observed_boundary,
                      "boundary_error": observed_boundary-predicted_boundary,
                      "rows": rows}, indent=2))

if __name__ == "__main__":
    main()