import json, random, numpy as np, torch
from experiment import AttnRegressor, make_data

def one(seed, device):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    tr_x,tr_y=make_data(256); te_x,te_y=make_data(64)
    tr_x,tr_y,te_x,te_y=[z.to(device) for z in (tr_x,tr_y,te_x,te_y)]
    out={}
    for biased in (False,True):
        torch.manual_seed(seed); m=AttnRegressor(biased).to(device)
        opt=torch.optim.Adam(m.parameters(),lr=3e-3); mse=torch.nn.MSELoss()
        for _ in range(180):
            ix=torch.randint(0,256,(16,),device=device)
            loss=mse(m(tr_x[ix]),tr_y[ix]); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=m(te_x); err=(pred-te_y).square().squeeze(-1)
            r=(te_x.square().sum(-1)).sqrt()
            near=err[r<.35].mean().item(); far=err[r>.7].mean().item()
            out['idea' if biased else 'baseline']=[mse(pred,te_y).item(),near,far]
    return out

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        vals=[one(s,device) for s in (17,23,41)]
    except Exception:
        device='cpu'; vals=[one(s,device) for s in (17,23,41)]
    print(json.dumps({'device':device,'seeds':[17,23,41],'columns':['overall_mse','near_origin_mse','far_mse'],'runs':vals},indent=2))
if __name__=='__main__': main()