import json, random, time
import numpy as np
import torch
from torch import nn
from third_order_optimizer import ThirdOrderLangevin

SEED=123

def seed(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def make_data(n=2048):
    g=np.random.default_rng(77)
    x=g.normal(size=(n,2)).astype('float32')
    y=((x[:,0]*x[:,1] + .25*x[:,0] - .15*x[:,1])>0).astype('int64')
    return torch.from_numpy(x), torch.from_numpy(y)

def run(kind, device):
    seed(SEED); x,y=make_data(); x,y=x.to(device),y.to(device)
    model=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2)).to(device)
    if kind=='sgd': opt=torch.optim.SGD(model.parameters(),lr=.08)
    elif kind=='momentum': opt=torch.optim.SGD(model.parameters(),lr=.04,momentum=.9)
    else: opt=ThirdOrderLangevin(model.parameters(),dt=.035,gamma=1.0,temperature=1e-4)
    lossfn=nn.CrossEntropyLoss(); losses=[]; t0=time.perf_counter()
    for step in range(300):
        opt.zero_grad(set_to_none=True); loss=lossfn(model(x),y); loss.backward(); opt.step()
        if step%50==0: losses.append(float(loss.detach().cpu()))
    with torch.no_grad(): acc=float((model(x).argmax(1)==y).float().mean().cpu())
    return {'loss_checkpoints':losses,'final_loss':float(loss.detach().cpu()),'train_accuracy':acc,'seconds':time.perf_counter()-t0}

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        out={k:run(k,device) for k in ['sgd','momentum','third_order']}
    except Exception:
        device='cpu'; out={k:run(k,device) for k in ['sgd','momentum','third_order']}
    result={'device':device,'steps':300,'results':out}
    with open('benchmark_results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()