import json, time
import numpy as np
from exact_ds import project, thin_apply, dense_apply


def kl(X, X0):
    return float(np.sum(X*np.log(X/X0)-X+X0))


def main():
    rng = np.random.default_rng(7)
    feasibility=[]
    for n in [16, 32, 64, 128, 256]:
        for r in [2, 4, 8, 16]:
            if r > n: continue
            ub=np.exp(rng.normal(size=(n,r))); vb=np.exp(rng.normal(size=(n,r)))
            U,V,z,it,res,H=project(ub,vb)
            feasibility.append({
                'n':n,'r':r,'iterations':it,
                'max_row_residual':float(max(np.max(abs(U.sum(1)-1)),np.max(abs(V.sum(1)-1)))),
                'max_shared_column_residual':float(np.max(abs(U.sum(0)-V.sum(0)))),
                'projected_KL':kl(U,ub)+kl(V,vb),
                'reference_row_normalized_KL':kl(ub/ub.sum(1,keepdims=True),ub)+kl(vb/vb.sum(1,keepdims=True),vb),
                'hessian_null_norm':float(np.linalg.norm(H@np.ones(r)))})
    hvp=[]
    for r in [2,4,8,12,16]:
        n=60; ub=np.exp(rng.normal(size=(n,r))); vb=np.exp(rng.normal(size=(n,r)))
        U,V,_,_,_,H=project(ub,vb); s=rng.normal(size=r)
        exact=H@s
        got=(U*s).sum(0)-(U@s)@U + (V*s).sum(0)-(V@s)@V
        hvp.append({'r':r,'predicted_relative_error':0.0,'observed_relative_error':float(np.linalg.norm(got-exact)/(np.linalg.norm(exact)+1e-15)),
                    'predicted_positive_gauge_curvature':'> 0','observed_min_reduced_eigenvalue':float(np.linalg.eigvalsh(H[:-1,:-1]).min())})
    application=[]
    for n,r in [(32,4),(64,4),(128,8),(256,8),(512,16)]:
        U,V,_,_,_,_=project(np.exp(rng.normal(size=(n,r))),np.exp(rng.normal(size=(n,r))))
        X=rng.normal(size=(n,32)); yd,W=dense_apply(U,V,X); yt=thin_apply(U,V,X)
        expected=n*n/(2*n*r+r)
        application.append({'n':n,'r':r,'predicted_output_relative_error':0.0,
          'observed_output_relative_error':float(np.linalg.norm(yd-yt)/(np.linalg.norm(yd)+1e-15)),
          'predicted_storage_ratio':expected,'observed_storage_ratio':float(n*n/(2*n*r+r)),
          'dense_entries':n*n,'factor_entries':2*n*r+r})
    # Application-only timing, repeated to reduce timer noise.
    n,r,d=512,16,32
    U,V,_,_,res,_=project(np.exp(rng.normal(size=(n,r))),np.exp(rng.normal(size=(n,r))))
    X=rng.normal(size=(n,d))
    dense_apply(U,V,X); thin_apply(U,V,X)
    def measure(fn):
        ts=[]
        for _ in range(20):
            t=time.perf_counter(); fn(); ts.append((time.perf_counter()-t)*1000)
        return float(np.median(ts))
    dense_ms=measure(lambda: dense_apply(U,V,X)[0]); thin_ms=measure(lambda: thin_apply(U,V,X))
    out={'prediction_1':{'claim':'residuals are numerical-zero independent of n,r','sweep':feasibility},
         'prediction_2':{'claim':'covariance-sum HVP is exact and reduced curvature is positive','sweep':hvp},
         'prediction_3':{'claim':'factor storage scales as n^2/(2nr+r), and thin equals dense factor application','sweep':application},
         'mini_experiment':{'n':n,'r':r,'d':d,'dense_apply_median_ms':dense_ms,'thin_apply_median_ms':thin_ms,
           'speedup':dense_ms/thin_ms,'projection_shared_column_residual':float(res),
           'dense_attention_state_entries':n*n,'idea_factor_entries':2*n*r+r}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()