import json
from experiment import train, evaluate, device, orbit_tangent, projector, make_data
import torch

def cond_stats():
    x,_=make_data(2000,909)
    R=orbit_tangent(x).unsqueeze(-1)
    # damped Gram condition, including singular points at the origin
    gram=R.transpose(1,2)@R + 1e-6*torch.eye(1)
    vals=gram[:,0,0]
    return {'min_damped_gram_eigenvalue':float(vals.min()),'max_damped_gram_eigenvalue':float(vals.max()),'condition_number':float(vals.max()/vals.min())}

rows=[]
for seed in [17,29,41]:
    b=evaluate(train(False,seed,350),seed+1000)
    q=evaluate(train(True,seed,350),seed+1000)
    rows.append({'seed':seed,'baseline':b,'idea':q,'drift_ratio':q[2]/b[2]})
out={'device':device,'runs':rows,'conditioning':cond_stats()}
print(json.dumps(out,indent=2))
with open('multiseed_results.json','w') as f: json.dump(out,f,indent=2)