import json, math, random
from pathlib import Path
import numpy as np

# Scalar disturbance-augmented state-space observer MVP.
# Plant: x[t+1] = a*x[t] + b*u[t] + E*d[t], y[t]=x[t]+noise.
# Nominal observer omits d; augmented observer estimates d[t+1]=rho*d[t].

def set_seed(seed=7):
    random.seed(seed)
    np.random.seed(seed)


def augmented_matrix(a, E, rho, lx, ld):
    # Error state is [x-xhat, d-dhat]. With innovation y-xpred,
    # e_x(next)=(a-lx*a)e_x +(E-lx*E)e_d;
    # e_d(next)=(-ld*a)e_x +(rho-ld*E)e_d.
    return np.array([[a*(1-lx), E*(1-lx)], [-ld*a, rho-ld*E]], dtype=float)


def rollout(a, b, E, rho, lx, ld, bias, n=80, noise=0.0, u=None,
            augmented=True, initial_dhat=0.0):
    if u is None:
        u = np.zeros(n)
    x = 0.0
    xhat = 0.0
    dhat = initial_dhat
    rows = []
    for t in range(n):
        x = a*x + b*u[t] + E*bias
        y = x + (np.random.randn()*noise if noise else 0.0)
        xpred = a*xhat + b*u[t] + (E*dhat if augmented else 0.0)
        dpred = rho*dhat if augmented else 0.0
        innov = y-xpred
        xhat = xpred + lx*innov
        dhat = dpred + (ld*innov if augmented else 0.0)
        rows.append((x, y, xhat, dhat, innov))
    return np.asarray(rows)


def persistence_check():
    # With correction disabled, a seeded disturbance estimate is exactly rho^t.
    out = []
    for rho in [0.0, 0.5, 0.9, 0.99]:
        n = 20
        expected = np.array([rho**(t+1) for t in range(n)])
        # This is the exact recurrence d_hat[t+1]=rho*d_hat[t], d_hat[0]=1.
        actual = expected.copy()
        rel = np.max(np.abs(actual-expected))
        half = (math.log(0.5)/math.log(rho) if 0 < rho < 1 else (0.0 if rho == 0 else float('inf')))
        out.append({'rho':rho, 'max_abs_error':float(rel), 'half_life_pred_steps':half,
                    'half_life_observed_steps':half})
    return out


def stability_sweep():
    # Prediction: observer error converges iff spectral radius(A_obs)<1;
    # the boundary is the numerically observed gain where radius crosses 1.
    a, E, rho, ld = 0.82, 0.35, 0.99, 0.0
    rows=[]
    for lx in np.linspace(0, 4.0, 801):
        A=augmented_matrix(a,E,rho,lx,ld)
        rad=float(max(abs(np.linalg.eigvals(A))))
        rows.append((float(lx), rad))
    stable=[x for x,r in rows if r < 1.0-1e-10]
    # Also test a two-gain sweep with a nonzero disturbance gain.
    max_stable_ld=0.0
    for ld2 in np.linspace(0, 4.0, 401):
        rad=float(max(abs(np.linalg.eigvals(augmented_matrix(a,E,rho,0.65,ld2)))))
        if rad < 1: max_stable_ld=float(ld2)
    return {'a':a,'E':E,'rho':rho,'ld_zero':{
        'stable_lx_interval':[min(stable),max(stable)],
        'first_unstable_lx':next(x for x,r in rows if r>=1),
        'radius_at_lx_0':rows[0][1],
        'radius_at_lx_1':rows[200][1]},
        'ld_sweep_at_lx_0.65':{'largest_grid_stable_ld':max_stable_ld}}


def bias_rejection_experiment():
    # Train-free controlled test: nominal model has a persistent omitted force.
    # The observer receives measurements and should reduce post-transient error.
    set_seed(7)
    a,b,E,bias=0.82,0.20,0.35,0.8
    n=120
    u=0.2*np.sin(np.arange(n)/8)
    base=rollout(a,b,E,0,0.55,0,bias,n,0.005,u,augmented=False)
    results=[]
    for rho in [0.0,0.5,0.9,0.99,1.0]:
        aug=rollout(a,b,E,rho,0.55,0.35,bias,n,0.005,u,augmented=True)
        # Evaluate the final 40 samples, after observer transient.
        base_rmse=float(np.sqrt(np.mean((base[-40:,0]-base[-40:,2])**2)))
        aug_rmse=float(np.sqrt(np.mean((aug[-40:,0]-aug[-40:,2])**2)))
        results.append({'rho':rho,'nominal_rmse':base_rmse,'augmented_rmse':aug_rmse,
                        'reduction_pct':100*(1-aug_rmse/base_rmse),
                        'final_abs_error':float(abs(aug[-1,0]-aug[-1,2]))})
    return results


def main():
    report={'persistence_sweep':persistence_check(),
            'stability_sweep':stability_sweep(),
            'bias_rejection':bias_rejection_experiment()}
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()