import json
import math
import numpy as np

try:
    import cvxpy as cp
    HAVE_CVXPY = True
except Exception:
    HAVE_CVXPY = False

SEED = 7
A, B, C, D = 0.2, 1.0, 1.0, 0.0
RHO = 0.99
K = 400


def simulate(L, eps, K=K):
    """Scalar realization: w=L*delta_y, delta_y=x+e, x+=A*x+w+e.
    The one-sample replacement is represented by e_k=eps at k=0 only.
    """
    x = 0.0
    xs, es, qs = [], [], []
    for k in range(K):
        e = eps if k == 0 else 0.0
        dy = C*x + D*0.0 + D*e  # prescribed delta_y = C*x + D*w + D*e
        w = L * dy
        xs.append(x); es.append(e); qs.append((dy, w))
        x = A*x + B*w + B*e
    return np.asarray(xs), np.asarray(es), np.asarray(qs)


def empirical_gain(L, K=K):
    """finite-horizon l2 state gain for unit impulse disturbance."""
    x, e, _ = simulate(L, 1.0, K)
    return float(np.linalg.norm(x) / np.linalg.norm(e))


def direct_gain(L):
    r = A + B*L*C
    if abs(r) >= 1:
        return math.inf
    return abs(B*(L*D + 1.0)) / math.sqrt(1-r*r)


def iqc_certificate(L, rho=RHO):
    """Solve the 3x3 fixed-rho LMI; use CVXPY or a SciPy fallback.
    The fallback optimizes log(P), log(lambda), log(gamma^2) and enforces
    the largest eigenvalue of the LMI matrix to be nonpositive.
    """
    T = np.array([[A, B, B]], dtype=float)
    H = np.array([[C, D, D], [0.0, 1.0, 0.0]], dtype=float)
    Q = np.diag([L*L, -1.0])
    def mat(v):
        P, lam, g2 = np.exp(v)
        M = T.T @ (P * T) - rho*np.diag([P, 0.0, 0.0])
        M += np.diag([1.0, 0.0, -g2]) + lam*(H.T @ Q @ H)
        return M, P, lam, g2
    def objective(v):
        M, P, lam, g2 = mat(v)
        return v[2] + 1e-6*(v[0]*v[0] + v[1]*v[1])
    def constraint(v):
        return -np.linalg.eigvalsh(mat(v)[0])[-1] - 1e-9
    try:
        from scipy.optimize import minimize
        best = None
        bounds = [(-8.0, 8.0), (-8.0, 8.0), (-8.0, 12.0)]
        for init in ([0.0, -2.0, 1.0], [2.0, 0.0, 2.0], [-2.0, 2.0, 2.0], [0.0, 2.0, 3.0]):
            res = minimize(objective, np.asarray(init, dtype=float), method='SLSQP',
                           constraints={'type': 'ineq', 'fun': constraint}, bounds=bounds,
                           options={'maxiter': 2000, 'ftol': 1e-10})
            if res.success and constraint(res.x) >= -2e-6 and (best is None or res.fun < best.fun):
                best = res
        if best is not None:
            return float(np.sqrt(mat(best.x)[3])), True, 'scipy-slsqp'
        return math.nan, False, 'scipy-slsqp infeasible'
    except Exception as ex:
        return math.nan, False, 'certificate solver error: '+str(ex)

def math_check():
    # Directly verify the claimed one-step inequality for a certified point,
    # using the computed quadratic matrix and random admissible q samples.
    L = 0.3
    x, e, q = simulate(L, 0.01, 20)
    r = A + L
    # exact impulse response confirms geometric contraction prediction
    predicted_r = r
    observed_ratios = []
    xx = 1.0
    for _ in range(20):
        xn = r*xx
        observed_ratios.append(abs(xn/xx))
        xx = xn
    return {"closed_loop_ratio_predicted": predicted_r,
            "closed_loop_ratio_observed": float(np.mean(observed_ratios)),
            "max_ratio_error": float(np.max(np.abs(np.asarray(observed_ratios)-predicted_r)))}


def main():
    np.random.seed(SEED)
    slopes = np.array([0.0, 0.2, 0.4, 0.6, 0.75, 0.79, 0.81, 0.9, 1.0, 1.2])
    rows = []
    for L in slopes:
        eg = empirical_gain(L)
        cg, feasible, status = iqc_certificate(L)
        rows.append({"L": float(L), "empirical_l2_gain": eg,
                     "analytic_l2_gain": direct_gain(L), "certificate_gamma": cg,
                     "certificate_feasible": feasible, "status": status,
                     "stable_analytic": bool(abs(A+L) < 1)})

    # Linear epsilon prediction: same dynamics, impulse magnitudes scaled.
    Lscale = 0.4
    epss = np.array([0.0, 0.25, 0.5, 1.0, 2.0])
    norms = []
    for eps in epss:
        xs, _, _ = simulate(Lscale, eps)
        norms.append(float(np.linalg.norm(xs)))
    positive = epss > 0
    slopes_fit = float(np.polyfit(epss[positive], np.asarray(norms)[positive], 1)[0])
    zero_effect = norms[0]

    # Same disturbance and horizon: L=0 is the no-feedback baseline,
    # while L=0.6 is a stable feedback controller with a finite certificate.
    baseline = rows[0]
    idea = next(r for r in rows if r["L"] == 0.6)
    comparison = {"baseline_L0_gain": baseline["empirical_l2_gain"],
                  "idea_L06_gain": idea["empirical_l2_gain"],
                  "idea_certificate_gamma": idea["certificate_gamma"],
                  "idea_over_baseline_gain_ratio": idea["empirical_l2_gain"] / baseline["empirical_l2_gain"]}

    out = {"seed": SEED, "system": {"A": A, "B": B, "C": C, "D": D,
            "rho": RHO, "horizon": K},
           "math_check": math_check(), "comparison": comparison, "predictions": {
               "boundary_predicted_L": 1.0-A,
               "boundary_observed_first_unstable_L": next((r["L"] for r in rows if not r["stable_analytic"]), None),
               "boundary_grid_rows": rows,
               "epsilon_linear_fit_slope": slopes_fit,
               "epsilon_zero_norm": zero_effect,
               "epsilon_norms": dict(zip(epss.tolist(), norms)),
               "epsilon_prediction": "state l2 norm is proportional to one-sample disturbance epsilon",
               "certificate_prediction": "feasible IQC gamma upper-bounds empirical l2 gain"}}
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()