"""MVP verification of a structured-mu feedback optimizer on a scalar quadratic.

The local error dynamics are x[t+1] = A x[t] + w[t], z[t] = r x[t],
where A=a-c is the nominal closed-loop pole, c is the optimizer feedback
(g_t=a*x_t, update u=-c*x), and w=delta*z, |delta|<=1.  Thus the uncertain
pole is A+delta*r.  The transfer from w to z is M(z)=r/(z-A), and the
scalar structured singular value is |M|.  Its frequency maximum is
r/(1-|A|), so the certificate is mu<1 iff |A|+r<1.
"""
import json
import math
import numpy as np
from dataclasses import dataclass

@dataclass
class StructuredMuController:
    """Low-order (zero-state) controller with a conservative sampled mu guard."""
    nominal_curvature: float
    uncertainty_radius: float
    target_mu: float = 0.90

    def choose_gain(self):
        # c=a makes A=0 and is the minimizer of the exact scalar certificate.
        # If the uncertainty itself is too large, no static controller can certify it.
        c = self.nominal_curvature
        achieved = self.mu_estimate(c)
        if achieved >= self.target_mu:
            # Keep the nominal pole at zero and expose an honest failure.
            return c, achieved, False
        return c, achieved, True

    def mu_estimate(self, c, nfreq=4097):
        A = self.nominal_curvature - c
        omega = np.linspace(0, 2*np.pi, nfreq)
        M = self.uncertainty_radius / (np.exp(1j*omega) - A)
        return float(np.max(np.abs(M)))

def exact_mu(a, r, c):
    A = a-c
    return r/(1-abs(A)) if abs(A) < 1 else float("inf")

def simulate(a, r, c, delta, steps=120, x0=1.0):
    x = float(x0)
    traj = [x]
    for _ in range(steps):
        x = (a-c+delta*r)*x
        traj.append(x)
        if not np.isfinite(x) or abs(x) > 1e12:
            break
    return np.asarray(traj)

def numerical_core_check():
    # Prediction 1: sampled frequency max equals r/(1-|A|), including its worst case.
    rows=[]
    for a,c,r in [(0.8,0.35,0.2),(0.8,0.8,0.65),(0.4,0.1,0.3)]:
        A=a-c
        w=np.linspace(0,2*np.pi,200001)
        sampled=float(np.max(np.abs(r/(np.exp(1j*w)-A))))
        analytic=exact_mu(a,r,c)
        rows.append({"a":a,"c":c,"r":r,"sampled_mu":sampled,"analytic_mu":analytic,
                     "relative_error":abs(sampled-analytic)/analytic})
    return rows

def boundary_sweep():
    # Prediction 2: for fixed A, transition is r*=1-|A|; measure actual trajectories.
    a,c=0.8,0.35
    A=a-c
    predicted=1-abs(A)
    rs=np.linspace(predicted-0.18,predicted+0.18,19)
    records=[]
    for r in rs:
        # delta=+1 is the worst structured perturbation for positive A.
        x=simulate(a,r,c,delta=1.0,steps=100)
        stable=bool(np.max(np.abs(x)) < 1e6 and abs(x[-1]) < abs(x[0]))
        records.append((float(r),stable,float(abs(x[-1]))))
    stable_rs=[r for r,s,_ in records if s]
    unstable_rs=[r for r,s,_ in records if not s]
    observed=(max(stable_rs)+min(unstable_rs))/2 if stable_rs and unstable_rs else float("nan")
    return {"A":A,"predicted_r_boundary":predicted,"observed_r_boundary":observed,
            "records":records,"boundary_error":abs(observed-predicted)}

def gain_sweep():
    # Prediction 3: c=a (zero nominal pole) minimizes worst-case contraction and mu.
    a,r=0.8,0.45
    gains=np.linspace(0.0,1.6,161)
    mus=np.array([exact_mu(a,r,c) for c in gains])
    contractions=np.array([abs(a-c)+r for c in gains])
    best_mu_gain=float(gains[np.argmin(mus)])
    best_contraction_gain=float(gains[np.argmin(contractions)])
    robust=StructuredMuController(a,r)
    c,mu,cert=robust.choose_gain()
    # baseline is a common under-corrected fixed step c=0.35
    base_c=0.35
    base_worst=abs(a-base_c)+r
    idea_worst=abs(a-c)+r
    return {"a":a,"r":r,"predicted_optimal_gain":a,
            "observed_mu_minimizer":best_mu_gain,
            "observed_contraction_minimizer":best_contraction_gain,
            "controller_gain":c,"controller_mu":mu,"certificate":cert,
            "baseline_gain":base_c,"baseline_worst_pole":base_worst,
            "idea_worst_pole":idea_worst,
            "mu_at_zero_uncertainty":exact_mu(a,0.0,a)}

def divergence_probability_sweep(seed=11):
    # Prediction: random bounded structured perturbations become unstable when
    # the worst-case certificate mu=r/(1-|A|) exceeds one.
    rng=np.random.default_rng(seed)
    a,c=0.8,0.35
    A=a-c
    rs=np.linspace(0.40,0.70,16)
    rows=[]
    for r in rs:
        divergent=0
        trials=1000
        for _ in range(trials):
            delta=float(rng.uniform(-1,1))
            pole=abs(A+delta*r)
            # Long-horizon growth from x0=1, with a conservative numerical cutoff.
            x=abs(pole)**80
            divergent += int(x>10)
        rows.append({"r":float(r),"mu":float(exact_mu(a,r,c)),
                     "divergence_probability":divergent/trials})
    return {"predicted_mu_boundary":1.0,"rows":rows}

def quadratic_mini_experiment(seed=7):
    # Same tiny setup, adversarial bounded curvature/noise represented by delta.
    rng=np.random.default_rng(seed)
    a=0.8; r=0.45; horizon=80; trials=200
    baseline_c=0.35
    idea=StructuredMuController(a,r); idea_c,_,_=idea.choose_gain()
    outcomes={"baseline":[],"idea":[]}
    for _ in range(trials):
        # Fixed but unknown structured perturbation plus small gradient noise.
        delta=float(rng.uniform(-1,1))
        noise=float(rng.normal(0,0.01))
        for name,c in [("baseline",baseline_c),("idea",idea_c)]:
            x=1.0; maxabs=1.0
            for _ in range(horizon):
                x=(a-c+delta*r)*x + noise
                maxabs=max(maxabs,abs(x))
            outcomes[name].append((abs(x),maxabs))
    summary={}
    for k,v in outcomes.items():
        v=np.asarray(v)
        summary[k]={"median_final_abs":float(np.median(v[:,0])),
                    "median_max_abs":float(np.median(v[:,1])),
                    "fraction_max_over_10":float(np.mean(v[:,1]>10))}
    summary["controller_gain"]=idea_c
    return summary

def main():
    out={"core_math":numerical_core_check(),"boundary":boundary_sweep(),
         "gain_prediction":gain_sweep(),"divergence_sweep":divergence_probability_sweep(),"mini_experiment":quadratic_mini_experiment()}
    print(json.dumps(out,indent=2))

if __name__ == "__main__": main()