import json
import numpy as np
from pathlib import Path

RNG = np.random.default_rng(1623)

# Damped oscillator: eigenvalues -alpha +/- i beta.  The known relaxation is -alpha*x;
# the sparse library learns the coupling terms [x1, x2].
alpha, beta, h = 0.20, 0.90, 0.08
M = np.array([[-alpha, -beta], [beta, -alpha]], dtype=float)


def trajectory(x0, n=260, noise=0.0):
    x = np.zeros((n, 2)); x[0] = x0
    D = np.eye(2) + h*M
    for k in range(n-1): x[k+1] = D @ x[k]
    if noise: x += RNG.normal(0, noise, x.shape)
    return x


def make_data(ntraj=24, n=180, obs_noise=0.003, outlier_frac=0.05):
    X, Y = [], []
    for _ in range(ntraj):
        clean = trajectory(RNG.normal(0, 1, 2), n, 0)
        obs = clean + RNG.normal(0, obs_noise, clean.shape)
        # centered derivative, with the same observation corruption mechanism
        der = (obs[2:] - obs[:-2])/(2*h)
        lib = obs[1:-1]
        # Known physics: f_known=-alpha*x. Learn only cross-state interactions.
        A_lib = np.column_stack([lib[:, 1], lib[:, 0]])
        residual_target = der + alpha*lib
        X.append(A_lib); Y.append(residual_target)
    A, b = np.vstack(X), np.vstack(Y)
    bad = RNG.choice(len(A), int(outlier_frac*len(A)), replace=False)
    b[bad] += RNG.normal(0, 0.7, (len(bad),2))
    return A, b, bad


def ols(A,b):
    return np.linalg.lstsq(A, b, rcond=None)[0]


def tls(A,b):
    # Total least squares independently for each derivative coordinate.
    out = np.zeros((A.shape[1], b.shape[1]))
    Z = np.column_stack([A, b])
    for j in range(b.shape[1]):
        # For multivariate b, the standard TLS null vector is applied to [A,b_j].
        _, _, vh = np.linalg.svd(np.column_stack([A, b[:,j]]), full_matrices=False)
        v = vh[-1]
        out[:,j] = -v[:-1]/v[-1]
    return out


def robust_tls(A,b, seed=4, rounds=500, threshold=0.08):
    rng = np.random.default_rng(seed); p=A.shape[1]
    best_in = None; best_score=-1
    m=min(len(A), p+3)
    for _ in range(rounds):
        idx=rng.choice(len(A),m,replace=False)
        try: c=tls(A[idx],b[idx])
        except np.linalg.LinAlgError: continue
        residual=np.linalg.norm(b-A@c,axis=1)
        # normalized residual makes threshold comparable across trials
        inliers=residual < threshold
        score=int(inliers.sum())
        if score>best_score: best_score, best_in=score,inliers
    if best_in is None: return tls(A,b), np.ones(len(A),bool)
    # TLS consensus refit, followed by iterative hard thresholding (sparsity step).
    c=tls(A[best_in],b[best_in])
    scale=np.maximum(np.max(np.abs(c),axis=0,keepdims=True),1e-9)
    c[np.abs(c)<0.08*scale]=0
    return c,best_in


def spectral_gain_boundary(M, hh):
    # J(g)=I+g*h*M; solve rho(J)=1 for the first positive crossing.
    eig=np.linalg.eigvals(M)
    candidates=[]
    for z in eig:
        if abs(z.imag)>1e-10:
            candidates.append(-2*z.real/(hh*abs(z)**2))
        elif z.real<0: candidates.append(-2/(hh*z.real))
    return min(candidates)


def rho(M, gain): return max(abs(np.linalg.eigvals(np.eye(2)+gain*h*M)))


def rollout_error(C, gain=1.0, n=100):
    # C maps row features to derivative targets; state-column dynamics use C.T.
    D=np.eye(2)+gain*h*C
    x=np.array([1.,-.4]); norms=[]
    for _ in range(n):
        norms.append(np.linalg.norm(x)); x=D@x
    return float(norms[-1]), float(max(norms))


def main():
    # Math sanity: exact boundary and a sweep around it.
    pred=spectral_gain_boundary(M,h)
    gain_grid=np.linspace(0.25, 1.35*pred, 45)
    measured=[]
    for g in gain_grid:
        # observed boundedness over 1000 steps, with a tiny nonzero initial state
        D=np.eye(2)+g*h*M; x=np.array([1.,0.]);
        for _ in range(1000): x=D@x
        measured.append(np.linalg.norm(x)<1e6)
    crossing=next((gain_grid[i] for i,v in enumerate(measured) if not v), gain_grid[-1])
    # Refine empirical boundary by direct spectral test on a dense grid.
    dense=np.linspace(0.1,1.3*pred,5000)
    obs=dense[np.argmax([rho(M,g)>=1 for g in dense])]

    A,b,bad=make_data()
    c_ols=ols(A,b); c_tls=tls(A,b); c_rob,inliers=robust_tls(A,b)
    q=np.array([[-beta,0.0],[0.0,beta]])
    errs={"OLS":float(np.linalg.norm(c_ols-q)/np.linalg.norm(q)),
          "TLS":float(np.linalg.norm(c_tls-q)/np.linalg.norm(q)),
          "TLS_RANSAC":float(np.linalg.norm(c_rob-q)/np.linalg.norm(q))}
    # Held-out clean 100-step trajectory one-step model rollout.
    def hybrid_matrix(c):
        # A=[x2,x1], so the learned interaction matrix is reconstructed here.
        return np.array([[-alpha, c[0,0]], [c[1,1], -alpha]])
    roll={}
    for name,c in [("OLS",c_ols),("TLS_RANSAC",c_rob)]: roll[name]=rollout_error(hybrid_matrix(c))
    roll["true"]=rollout_error(M)

    # Step-size scaling prediction: g* = 2 alpha/(h(alpha^2+beta^2)), measured by eigenvalue sweep.
    scaling=[]
    for hh in [0.04,0.08,0.16,0.24]:
        p=spectral_gain_boundary(M,hh)
        gs=np.linspace(.1,1.15*p,4000)
        measured_h=gs[np.argmax([max(abs(np.linalg.eigvals(np.eye(2)+g*hh*M)))>=1 for g in gs])]
        scaling.append({"h":hh,"predicted":p,"observed":float(measured_h),"relative_error":float(abs(measured_h-p)/p)})
    result={"parameters":{"alpha":alpha,"beta":beta,"h":h,"outlier_fraction":len(bad)/len(A)},
      "stability_boundary":{"predicted_gain":pred,"observed_gain":float(obs),"coarse_rollout_crossing":float(crossing),"relative_error":float(abs(obs-pred)/pred),"stable_at_0.9":bool(rho(M,.9*pred)<1),"unstable_at_1.1":bool(rho(M,1.1*pred)>1)},
      "step_scaling":scaling,
      "coefficient_relative_error":errs,
      "inlier_fraction":float(inliers.mean()),"rollout_norm_100":roll}
    Path("results.json").write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=="__main__": main()