import json
from pathlib import Path
import numpy as np
from scipy.linalg import eig


def pencil(y,n,s=0):
    H0=np.array([[y[s+i+j] for j in range(n)] for i in range(n)],float)
    H1=np.array([[y[s+i+j+1] for j in range(n)] for i in range(n)],float)
    u,sv,vh=np.linalg.svd(H0); r=max(1,int(np.sum(sv>1e-9*sv[0])))
    U=u[:,:r]; V=vh.T[:,:r]
    z=eig(U.T@H1@V,U.T@H0@V,right=False)
    return z[np.isfinite(z)]


def winding(values):
    ph=np.unwrap(np.angle(values))
    return int(np.rint((ph[-1]-ph[0])/(2*np.pi)))


def certify(y,n=4,shifts=(0,3,6,9,12),tol=.10):
    rec=[]
    for s in shifts:
        z=pencil(y,n,s); rec.append(z[np.abs(z)<1.05])
    ref=rec[0]; chosen=[]
    for z in ref:
        support=sum(bool(len(a)) and np.min(np.abs(a-z))<tol for a in rec)
        if support>=max(2,len(shifts)//2): chosen.append(z)
    return np.asarray(chosen,complex),rec


def fit_predict(y,z,train):
    if len(z)==0: return np.inf
    tt=np.arange(train); X=np.stack([zj**tt for zj in z],1)
    a=np.linalg.lstsq(X,y[:train],rcond=None)[0]
    all_t=np.arange(len(y)); Xa=np.stack([zj**all_t for zj in z],1)
    return float(np.mean((np.real(Xa@a)[train:]-y[train:])**2))


def one_trial(noise,seed):
    rng=np.random.default_rng(seed); T=260; t=np.arange(T)
    true=np.array([.96*np.exp(1j*.31),.88*np.exp(1j*.73)])
    amp=np.array([1+.2j,.65-.1j]); clean=np.real(sum(a*z**t for a,z in zip(amp,true)))
    y=clean+noise*rng.normal(size=T)
    cert,rec=certify(y)
    # A real sequence requires conjugate completion; reject incomplete candidates.
    pos=[z for z in cert if z.imag>1e-7 and abs(z)<1]
    idea=np.concatenate([np.array(pos),np.conj(np.array(pos))]) if pos else np.array([],complex)
    base=np.array([.75*np.exp(1j*.2),.75*np.exp(1j*.6),.75*np.exp(1j*.9),.75*np.exp(1j*1.2)])
    persistence=[int(sum(bool(a.size) and np.min(np.abs(a-z))<.10 for a in rec)) for z in pos]
    return {"certified_pairs":len(pos),"persistence":persistence,
            "idea_mse":fit_predict(y,idea,100),"baseline_mse":fit_predict(y,base,100)}


def run():
    # Pure sequence verifies the generalized-eigenvalue claim to numerical precision.
    t=np.arange(40); true=np.array([.96*np.exp(1j*.31),.88*np.exp(1j*.73)])
    clean=np.real((1+.2j)*true[0]**t+(.65-.1j)*true[1]**t)
    z=pencil(clean,4,0)
    eig_err=max(float(np.min(np.abs(z-a))) for a in true)
    p=.8*np.exp(2j*np.pi*np.arange(512)/512)
    q=np.prod(1-true[:,None]*p[None,:],axis=0)
    math={"noiseless_max_eigenvalue_error":eig_err,"winding_radius_0.8":winding(q),"expected_winding":0,"contour_margin":float(np.min(np.abs(q)))}
    trials={}
    for noise in (.01,.035,.10):
        vals=[one_trial(noise,s) for s in (7,17,27,37,47)]
        trials[str(noise)]={"mean":{k:float(np.mean([v[k] if np.isscalar(v[k]) else np.mean(v[k]) for v in vals])) for k in ("certified_pairs","idea_mse","baseline_mse")},"trials":vals}
    out={"math_check":math,"noise_sweep":trials,"seed_set":[7,17,27,37,47]}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))

if __name__=='__main__': run()