import json
import numpy as np
from prolate_bottleneck import ProlateBottleneck, concentration_matrix, fourier_basis

rng=np.random.default_rng(2368)
T,W,delta=96,0.22,0.1
C=concentration_matrix(T,W)
e,U=np.linalg.eigh(C); o=np.argsort(e)[::-1]; e=e[o]; U=U[:,o]
# band-limited covariance samples
X=rng.standard_normal((1200,T)) @ (U*np.sqrt(np.maximum(e,0))).T
b=ProlateBottleneck(T,W,delta)
r=b.rank
Uf=fourier_basis(T,r)
Ur=np.linalg.qr(rng.standard_normal((T,r)))[0]
def err(Q):
    R=X-X@Q@Q.T
    return float(np.mean(np.sum(R*R,axis=1)/np.maximum(np.sum(X*X,axis=1),1e-12)))
# API check on [B,T,D]
X3=X[:4,:,None]
api=float(np.max(np.abs(b.reconstruct(X3)[:,:,0] - X3[:,:,0] @ b.basis @ b.basis.T)))
out={
 'rank':r,'empirical_rank':b.empirical_rank,'asymptotic_rank_ceiling':r,
 'eigen_min':float(e[-1]),'eigen_max':float(e[0]),'trace':float(np.trace(C)),
 'dpss_error':err(b.basis),'fourier_error':err(Uf),'random_error':err(Ur),
 'dpss_vs_fourier_relative_improvement':float((err(Uf)-err(b.basis))/err(Uf)),
 'quadratic_work_reduction':float(1-(r/T)**2),'api_max_abs_error':api,
 'curve':[]}
for q in [8,16,24,32,40,48]:
    Q=b.full_basis[:,:q]
    F=fourier_basis(T,q)
    R=np.linalg.qr(rng.standard_normal((T,q)))[0]
    out['curve'].append({'rank':q,'dpss':err(Q),'fourier':err(F),'random':err(R)})
open('validation_results.json','w').write(json.dumps(out,indent=2))
print(json.dumps(out,indent=2))