"""Matched ordinary/branch-conditioned MLP baselines. Grouped splits only."""
import os
os.environ.setdefault('OMP_NUM_THREADS','1')
from pathlib import Path
import json,time
import numpy as np
import torch
from torch import nn
from safetensors.torch import save_file
ROOT=Path(__file__).resolve().parents[1]
torch.set_num_threads(2)

def features(c,branch=None):
    chi=c['chi'];a=c['fx']+2*c['h']
    x=[np.log10(chi),a/chi,c['fy']/chi,c['moment']/chi,4*c['h']*c['delta']/chi,c['k1'],c['k2'],c['k3']]
    if branch is not None:x += [float(branch==b) for b in ['down','straight','up']]
    return x

def network(d):return nn.Sequential(nn.Linear(d,96),nn.Tanh(),nn.Linear(96,96),nn.Tanh(),nn.Linear(96,48))

def reconstruct(z,rows):
    full=np.column_stack([np.zeros(len(z)),z]);th=(full[:,1:]+full[:,:-1])/2
    a=np.array([v['controls']['fx']+2*v['controls']['h'] for v in rows])[:,None]
    fy=np.array([v['controls']['fy'] for v in rows])[:,None]
    stretch=1+a*np.cos(th)+fy*np.sin(th)
    xy=np.stack([np.cos(th),np.sin(th)],axis=-1)*stretch[:,:,None]/48
    xy=np.concatenate([np.zeros((len(z),1,2)),np.cumsum(xy,axis=1)],axis=1)
    return xy

def main():
    rows=[json.loads(l) for l in (ROOT/'data/equilibria.jsonl').open()]
    splits={s:np.array([i for i,r in enumerate(rows) if r['split']==s]) for s in ['train','validation','test','parameter_ood']}
    truth=np.array([r['z'] for r in rows],dtype='float32')
    out=ROOT/'models';out.mkdir(exist_ok=True);report={}
    for conditioned in [False,True]:
        name='branch_conditioned' if conditioned else 'ordinary'
        torch.manual_seed(14);rng=np.random.default_rng(14)
        x=np.array([features(r['controls'],r['branch'] if conditioned else None) for r in rows],dtype='float32')
        mean=x[splits['train']].mean(0);std=x[splits['train']].std(0);std[std<1e-6]=1
        X=torch.from_numpy((x-mean)/std);Y=torch.from_numpy(truth)
        net=network(x.shape[1]);opt=torch.optim.Adam(net.parameters(),lr=.002)
        best=float('inf');best_epoch=0;start=time.perf_counter()
        for epoch in range(401):
            net.train()
            for idx in np.array_split(rng.permutation(splits['train']),19):
                opt.zero_grad();loss=(net(X[idx])-Y[idx]).square().mean();loss.backward();opt.step()
            net.eval()
            with torch.no_grad():val=float((net(X[splits['validation']])-Y[splits['validation']]).square().mean())
            if val<best:
                best=val;best_epoch=epoch;weights={k:v.clone().contiguous() for k,v in net.state_dict().items()}
            if epoch%100==0:print(name,epoch,'val angle MSE',val,flush=True)
        net.load_state_dict(weights);save_file(weights,str(out/f'{name}.safetensors'))
        (out/f'{name}.json').write_text(json.dumps(dict(mean=mean.tolist(),std=std.tolist(),conditioned=conditioned,architecture=[x.shape[1],96,96,48],activation='tanh',n=48,features='see inference.py / features',best_epoch=best_epoch),indent=2))
        with torch.no_grad():pred=net(X).numpy()
        predicted=reconstruct(pred,rows);target=reconstruct(truth,rows)
        result={'training_seconds':time.perf_counter()-start,'best_validation_angle_mse':best,'best_epoch':best_epoch,'splits':{}}
        for s,idx in splits.items():
            dist=np.linalg.norm(predicted[idx]-target[idx],axis=2)
            br=np.where(np.max(np.abs(pred[idx]),axis=1)<.01,'straight',np.where(predicted[idx,-1,1]>0,'up','down'))
            result['splits'][s]=dict(states=len(idx),centerline_rmse_over_L=float(np.sqrt(np.mean(dist**2))),tip_mean_error_over_L=float(np.mean(dist[:,-1])),tip_p95_error_over_L=float(np.quantile(dist[:,-1],.95)),branch_accuracy=float(np.mean(br==np.array([rows[i]['branch'] for i in idx]))))
            result['splits'][s]['families']={}
            for family in ['end_load','intrinsic_curvature','microforce']:
                select=np.array([rows[i]['family']==family for i in idx])
                result['splits'][s]['families'][family]=dict(centerline_rmse_over_L=float(np.sqrt(np.mean(dist[select]**2))),tip_mean_error_over_L=float(np.mean(dist[select,-1])))
        report[name]=result
    (ROOT/'reports/baselines.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2))

if __name__=='__main__':main()