import json
import numpy as np
from scipy.linalg import eigvalsh

# Scalar latent plant: z_dot = A z, y = C z + w.
# The observation disturbance w enters the observer error through -K w.
A, C = -0.4, 1.0
alpha, gamma = 0.2, 0.9


def lmi_matrix(K, P):
    Ac = A - K * C
    Bc = -K                 # Bw=0, Dw=1
    Ce, De = 1.0, 0.0       # performance output q=e
    return np.array([[2 * Ac * P + Ce**2 + 2 * alpha * P,
                     P * Bc + Ce * De],
                     [P * Bc + Ce * De,
                      De**2 - gamma**2]])


def best_storage(K):
    # Search the one-dimensional positive storage metric, as a transparent
    # numerical analogue of optimizing P=LL^T+p_min I.
    ps = np.logspace(-4, 3, 600)
    vals = np.array([eigvalsh(lmi_matrix(K, p))[-1] for p in ps])
    i = int(vals.argmin())
    return float(ps[i]), float(vals[i])


def exact_error(K, e0, times, w):
    # Piecewise-constant forcing, exact propagation over each small interval.
    e = np.empty(len(times)); e[0] = e0
    dt = times[1] - times[0]
    Ac = A - K
    for j in range(len(times) - 1):
        # e_dot = Ac e - K w[j]
        e[j+1] = np.exp(Ac*dt)*e[j] + (-K*w[j]) * np.expm1(Ac*dt)/Ac
    return e


def dissipation_check(K, seed=7):
    rng = np.random.default_rng(seed)
    dt, T = 0.001, 8.0
    t = np.arange(0, T + dt, dt)
    # Smooth, piecewise disturbance, including an impulsive-like short pulse.
    w = 0.35 * rng.standard_normal(len(t)-1)
    w = np.convolve(w, np.ones(20)/20, mode='same')
    w[1000:1030] += 5.0
    e0 = 1.3
    e = exact_error(K, e0, t, w)
    P, lmax = best_storage(K)
    V = P * e**2
    lhs = (V[-1] - V[0] + np.sum((2*alpha*V[:-1] + e[:-1]**2 - gamma**2*w**2)*dt))
    # Direct induced-energy ratio (excluding initial-condition energy).
    zero = exact_error(K, 0.0, t, w)
    ratio = np.sum(zero[:-1]**2)*dt / (np.sum(w**2)*dt)
    return dict(P=P, lmi_max_eigenvalue=lmax, dissipation_residual=float(lhs),
                induced_energy_ratio=float(ratio), initial_V=float(V[0]),
                final_V=float(V[-1]))


def exact_hinf_gain(K):
    # For e_dot=-(A_abs+K)e-Kw, q=e, the stable transfer is -K/(s+0.4+K).
    # Its H-infinity norm is attained at zero frequency.
    return float(K / (K - A))


def zero_disturbance_decay(K):
    dt, T, e0 = 0.002, 8.0, 1.0
    t = np.arange(0, T+dt, dt)
    e = exact_error(K, e0, t, np.zeros(len(t)-1))
    # Compare observed decay exponent with required V decay exponent 2 alpha.
    P, _ = best_storage(K)
    V = P*e**2
    observed_rate = -np.log(V[-1]/V[0]) / T
    return dict(observed_V_decay_rate=float(observed_rate), required_rate=2*alpha,
                V_initial=float(V[0]), V_final=float(V[-1]))


def main():
    # Baseline: unconstrained high-gain innovation correction.
    baseline_K = 10.0
    # Idea: choose the largest gain on a coarse grid that is LMI feasible,
    # then refine by the storage search. This is a tiny proxy for training K,P.
    candidates = np.linspace(0.05, 10.0, 300)
    feasible = [(K, best_storage(K)) for K in candidates if best_storage(K)[1] <= 1e-8]
    idea_K, (idea_P, idea_lmi) = feasible[-1]
    result = {
        'config': {'A': A, 'C': C, 'alpha': alpha, 'gamma': gamma},
        'baseline': {'K': baseline_K, 'lmi': best_storage(baseline_K),
                     'exact_hinf_gain': exact_hinf_gain(baseline_K),
                     'zero_disturbance': zero_disturbance_decay(baseline_K),
                     'disturbance': dissipation_check(baseline_K)},
        'idea': {'K': float(idea_K), 'lmi': (idea_P, idea_lmi),
                 'exact_hinf_gain': exact_hinf_gain(idea_K),
                 'zero_disturbance': zero_disturbance_decay(idea_K),
                 'disturbance': dissipation_check(idea_K)},
        'claim_check': {
            'idea_lmi_feasible': bool(idea_lmi <= 1e-8),
            'baseline_lmi_feasible': bool(best_storage(baseline_K)[1] <= 1e-8),
            'idea_dissipation_holds': bool(dissipation_check(idea_K)['dissipation_residual'] <= 1e-5),
            'idea_exact_hinf_below_gamma': bool(exact_hinf_gain(idea_K) <= gamma),
            'baseline_exact_hinf_below_gamma': bool(exact_hinf_gain(baseline_K) <= gamma)
        }
    }
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()