"""Load safetensors baseline; optionally correct using the released mechanics."""
from pathlib import Path
import json
import numpy as np
import torch
from safetensors.torch import load_file
from train import features,network
from rod import Rod,Controls

def predict(controls,branch='down',model_dir=None,correct=False):
    folder=Path(model_dir) if model_dir else Path(__file__).resolve().parents[1]/'models'
    name='branch_conditioned' if branch is not None else 'ordinary'
    config=json.loads((folder/f'{name}.json').read_text())
    net=network(config['architecture'][0]);net.load_state_dict(load_file(str(folder/f'{name}.safetensors')));net.eval()
    x=(np.array(features(controls,branch))-config['mean'])/config['std']
    with torch.no_grad():z=net(torch.tensor(x,dtype=torch.float32)).numpy()
    rod=Rod(48);c=Controls(**controls)
    if correct:return rod.solve(c,z)
    th=rod.A@z;stretch=1+(c.fx+2*c.h)*np.cos(th)+c.fy*np.sin(th)
    xy=np.vstack([np.zeros(2),np.cumsum(stretch[:,None]*np.column_stack([np.cos(th),np.sin(th)])/48,axis=0)])
    return dict(z=z,centerline=xy,tip=xy[-1],note='Raw neural estimate; stability is not certified.')

if __name__=='__main__':
    c=dict(chi=.0001,fx=-.016428,fy=0.,moment=0.,h=.01,delta=.1,k1=0.,k2=0.,k3=0.)
    for corrected in [False,True]:
        r=predict(c,'down',correct=corrected)
        print('corrected' if corrected else 'neural','tip',r['tip'],'residual',r.get('residual'),'stability',r.get('stability'))