import json
from pathlib import Path
import numpy as np

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


def fit_dynamics(h, z, ridge=1e-6):
    # h[t+1] = A h[t] + B z[t] + e[t]
    X = np.concatenate([h[:-1], z], axis=1)
    Y = h[1:]
    W = Y.T @ X @ np.linalg.inv(X.T @ X + ridge*np.eye(X.shape[1]))
    d, p = h.shape[1], z.shape[1]
    return W[:, :d], W[:, d:], Y - X @ W.T


def q_factor(F, P):
    # Similarity-invariant generalized contraction factor.
    Pm = np.linalg.inv(np.linalg.cholesky(P))
    M = Pm @ F @ P @ F.T @ Pm.T
    return float(np.sqrt(max(np.linalg.eigvalsh(M).max(), 0)))


def scalar_sweep():
    # h+=q h + bounded disturbance. Exact worst-case radius is sigma/(1-q).
    qs = np.array([0.50, 0.70, 0.85, 0.95, 0.99, 1.01, 1.10])
    sigma = 0.1
    rows = []
    for q in qs:
        # adversarial disturbance maximizes absolute state and exposes the formula.
        h = 0.0
        vals = []
        for _ in range(600):
            e = sigma if h >= 0 else -sigma
            h = q*h + e
            vals.append(abs(h))
        if q < 1:
            predicted = sigma/(1-q)
            observed = float(np.mean(vals[-50:]))
            relerr = abs(observed-predicted)/predicted
        else:
            predicted = float('inf')
            observed = float(np.mean(vals[-20:]))
            relerr = None
        growth = (float(np.log(max(vals[-1], 1e-300)/max(vals[-21], 1e-300))/20) if q > 1 else None)
        predicted_growth = float(np.log(q)) if q > 1 else None
        rows.append(dict(q=float(q), predicted_radius=predicted,
                         observed_tail_radius=observed, relative_error=relerr,
                         diverged=bool(max(vals)>1e6), observed_log_growth_per_step=growth,
                         predicted_log_growth_per_step=predicted_growth))
    return rows


def boundary_sweep():
    # Random bounded disturbances: classify boundedness across a fine q sweep.
    qs = np.linspace(.85, 1.15, 31)
    sigma = .02
    rows=[]
    for q in qs:
        h=0.; maxabs=0.
        for _ in range(400):
            h=q*h+rng.uniform(-sigma,sigma)
            maxabs=max(maxabs,abs(h))
        rows.append((float(q), float(maxabs)))
    # operational divergence threshold is first q whose final state exceeds 10x stable q=.99 scale
    threshold = next((q for q,m in rows if m > 2.0), None)
    return rows, threshold


def data_fit_and_projection():
    # Offline trajectory fit, covariance ellipsoid, and same forcing with/without projection.
    d=3; p=2
    A=np.array([[.72,.08,0],[-.04,.64,.05],[0,.03,.58]])
    B=np.array([[.35,.0],[0,.28],[.05,.12]])
    T=5000
    z=rng.normal(size=(T,p))*0.3
    h=np.zeros((T+1,d))
    for t in range(T):
        h[t+1]=A@h[t]+B@z[t]+rng.normal(size=d)*.006
    Ah,Bh,res=fit_dynamics(h,z)
    # diagonal inflated residual ellipsoid (99.5% absolute coordinate containment)
    rad=np.quantile(np.abs(res), .995, axis=0)*1.05
    P=np.diag((rad/(1-.72))**2) # conservative axis box-to-ellipsoid scale
    fit_err=float(np.linalg.norm(A-Ah)/np.linalg.norm(A))
    # Deliberately unstable learned/recurrent matrix to make projection effect visible.
    F=Ah*1.55
    qp=q_factor(F,P)
    x=np.zeros(d); xp=np.zeros(d); raw=[]; proj=[]
    for _ in range(300):
        noise=rng.normal(size=d)*rad
        x=F@x+noise
        xp=F@xp+noise
        energy=float(xp@np.linalg.inv(P)@xp)
        if energy>1: xp=xp/np.sqrt(energy)
        raw.append(float(x@np.linalg.inv(P)@x)); proj.append(float(xp@np.linalg.inv(P)@xp))
    return {
        'fit_relative_A_error':fit_err, 'residual_axis_radii':rad.tolist(),
        'q_of_test_matrix':qp, 'raw_violation_rate':float(np.mean(np.array(raw)>1)),
        'projected_violation_rate':float(np.mean(np.array(proj)>1+1e-10)),
        'raw_max_ellipsoid_energy':float(max(raw)), 'projected_max_ellipsoid_energy':float(max(proj)),
        'projected_violation_tolerance':1e-10}


def main():
    radii=scalar_sweep()
    b_rows,boundary=boundary_sweep()
    fit=data_fit_and_projection()
    stable=[r for r in radii if r['q']<1]
    mean_err=float(np.mean([r['relative_error'] for r in stable]))
    report={
      'seed':SEED,
      'predictions':{
        'P1_boundary_predicted_q':1.0,
        'P1_observed_operational_boundary_q':boundary,
        'P2_radius_formula_mean_relative_error':mean_err,
        'P2_radius_sweep':radii,
        'P3_projection_predicted_zero_violation':True,
        'P3_projection_result':fit
      },
      'boundary_sweep': [{'q':q,'max_abs_state':m} for q,m in b_rows]
    }
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__=='__main__': main()