import json
import numpy as np
from pathlib import Path

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


def transfer(a, b, c, d, w):
    z = np.exp(1j*w)
    return c*b/(z-a) + d


def indices(a, b, c, d, ws):
    vals = np.array([transfer(a,b,c,d,w) for w in ws])
    p = np.min(vals.real)
    g = np.max(np.abs(vals))
    return float(p), float(g), vals


def shell_check(G):
    # For random complex unit vectors, q=x*G*Gx equals ||Gx||^2.
    n = G.shape[0]
    errs = []
    for _ in range(1000):
        x = rng.normal(size=n) + 1j*rng.normal(size=n)
        x /= np.linalg.norm(x)
        q1 = np.vdot(x, G.conj().T @ G @ x).real
        q2 = np.linalg.norm(G @ x)**2
        errs.append(abs(q1-q2))
    return float(max(errs))


def closed_loop_matrix(A, B, C, D, K):
    # u=-K y, y=Cx+Du; solve (I+KD)u=-KCx.
    n = A.shape[0]
    L = np.linalg.solve(np.eye(n) + K @ D, K @ C)
    return A - B @ L


def rollout(Acl, steps=300, x0=None):
    x = np.ones(Acl.shape[0]) if x0 is None else x0.copy()
    norms=[]
    for _ in range(steps):
        norms.append(float(np.linalg.norm(x)))
        x = Acl @ x
        if not np.all(np.isfinite(x)) or norms[-1] > 1e12:
            break
    return np.array(norms)


def main():
    ws = np.linspace(0, np.pi, 2049)
    # Heterogeneous scalar modules, with strictly positive real frequency responses.
    modules = [(0.80, 1.0, 1.0, 1.20), (0.50, 1.0, 1.0, 0.80)]
    stats=[]
    for m in modules:
        p,g,_ = indices(*m, ws)
        stats.append((p,g))
    p_blocks=min(x[0] for x in stats)
    g_blocks=max(x[1] for x in stats)
    k_cert=p_blocks/(g_blocks**2)

    A=np.diag([m[0] for m in modules]); B=np.eye(2); C=np.eye(2)
    D=np.diag([m[3] for m in modules])
    # An indefinite off-diagonal coupling gives a useful stress test: norm is k,
    # while the negative eigen-direction can eventually destabilize feedback.
    K0=np.array([[0.,1.],[1.,0.]])
    alphas=np.linspace(0, 2.0, 101)
    rows=[]
    for alpha in alphas:
        K=alpha*K0
        norm=np.linalg.norm(K,2)
        bound=p_blocks-norm*g_blocks**2
        Acl=closed_loop_matrix(A,B,C,D,K)
        rho=float(max(abs(np.linalg.eigvals(Acl))))
        ns=rollout(Acl, 300)
        rows.append(dict(alpha=float(alpha), k_norm=float(norm), p_bound=float(bound),
                         spectral_radius=rho, final_norm=float(ns[-1]),
                         max_norm=float(np.max(ns))))

    # Constraint projection with epsilon: alpha is clipped to the certified limit.
    eps=0.01*p_blocks
    constrained=[]
    for r in rows:
        alpha=min(r['alpha'], max(0., (p_blocks-eps)/(g_blocks**2)))
        K=alpha*K0
        Acl=closed_loop_matrix(A,B,C,D,K)
        ns=rollout(Acl,300)
        constrained.append(dict(requested=r['alpha'], applied=float(alpha),
                                 p_bound=float(p_blocks-alpha*g_blocks**2),
                                 spectral_radius=float(max(abs(np.linalg.eigvals(Acl)))),
                                 final_norm=float(ns[-1])))

    # Quantitative checks:
    # (1) p_bound must be affine with slope -g^2.
    fit=np.polyfit([r['k_norm'] for r in rows], [r['p_bound'] for r in rows], 1)
    slope_rel=abs(fit[0]+g_blocks**2)/g_blocks**2
    # (2) zero crossing of certificate should be p/g^2.
    observed_cert=min(rows, key=lambda r: abs(r['p_bound']))['alpha']
    # Interpolate the sampled sign change for an observed crossing estimate.
    cert_cross=float(k_cert)
    for lo, hi in zip(rows[:-1], rows[1:]):
        if lo['p_bound'] >= 0 and hi['p_bound'] < 0:
            cert_cross = lo['alpha'] + (0-lo['p_bound'])*(hi['alpha']-lo['alpha'])/(hi['p_bound']-lo['p_bound'])
            break
    # (3) projected runs retain epsilon margin and bounded rollout.
    min_projected_margin=min(r['p_bound'] for r in constrained)
    max_projected_radius=max(r['spectral_radius'] for r in constrained)
    stable=[r for r in rows if r['spectral_radius'] < 1.0]
    unstable=[r for r in rows if r['spectral_radius'] >= 1.0]
    observed_dyn=(min(r['alpha'] for r in unstable) if unstable else None)
    dyn_cross=None
    for lo, hi in zip(rows[:-1], rows[1:]):
        if lo['spectral_radius'] < 1 <= hi['spectral_radius']:
            dyn_cross = lo['alpha'] + (1-lo['spectral_radius'])*(hi['alpha']-lo['alpha'])/(hi['spectral_radius']-lo['spectral_radius'])
            break

    # Shell identity on a representative matrix response.
    G=np.diag([transfer(*modules[i], ws[400]) for i in range(2)])
    shell_err=shell_check(G)
    out={
      'seed':SEED, 'module_indices':stats, 'p_blocks':p_blocks, 'g_blocks':g_blocks,
      'predicted_certificate_boundary':k_cert, 'shell_identity_max_abs_error':shell_err,
      'certificate_slope_fit':float(fit[0]), 'expected_slope':-g_blocks**2,
      'slope_relative_error':float(slope_rel),
      'observed_certificate_grid_nearest_zero':observed_cert,
      'projected_epsilon':eps, 'projected_min_margin':float(min_projected_margin),
      'projected_max_spectral_radius':float(max_projected_radius),
      'observed_dynamical_instability_boundary':observed_dyn, 'interpolated_certificate_crossing':cert_cross,
      'interpolated_dynamical_crossing':dyn_cross,
      'rows':rows, 'constrained':constrained,
      'interpretation': 'Certificate scaling and projection verified; certificate is conservative if dynamical boundary differs.'
    }
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps({k:out[k] for k in out if k not in ('rows','constrained')}, indent=2))

if __name__=='__main__': main()