"""Fit local free-space response models; reserve 8 mrad steps for validation.

These are empirical command-to-measurement models, not UR firmware models or
contact dynamics. Cross-clock error is included in the fitted effective delay.
"""
import argparse
import json
from pathlib import Path

import numpy as np
from scipy.optimize import least_squares
from scipy.signal import cont2discrete, lfilter


def simulate(command_time, command, sample_time, params, initial=0., dt=.001):
    """Causal delay + first/second-order filter, internally sampled at 1 ms.

    ZOH the actual command sequence. Delay is applied to the continuous output
    approximation. Input/output units match; bias uses those same units.
    """
    grid = np.arange(min(0., sample_time[0], command_time[0])-.1,
                     max(sample_time[-1], command_time[-1])+.01, dt)
    idx = np.searchsorted(command_time, grid, side='right')-1
    u = command[np.clip(idx, 0, len(command)-1)]-initial
    u[idx < 0] = 0.
    if params['kind'] == 'first':
        num, den = [1.], [params['tau_s'], 1.]
    else:
        w, z = params['wn_rad_s'], params['zeta']
        num, den = [w*w], [1., 2*z*w, w*w]
    b, a, _ = cont2discrete((num, den), dt, method='zoh')
    y = lfilter(b.ravel(), a, u)
    return initial+params.get('gain',1.)*np.interp(sample_time-params['delay_s'],grid,y,left=0.)+params.get('bias',0.)


def fit_model(ct, command, t, measured, train, kind):
    def unpack(x):
        if kind == 'first':
            return dict(kind=kind,delay_s=x[0],tau_s=x[1],gain=x[2],bias=x[3])
        return dict(kind=kind,delay_s=x[0],wn_rad_s=x[1],zeta=x[2],gain=x[3],bias=x[4])
    if kind=='first':
        starts=[[d,.08,1.,0.] for d in [.0,.015,.03,.05]]
        bounds=([0.,.005,.8,-.3],[.12,.3,1.2,.3])
    else:
        starts=[[d,25.,z,1.,0.] for d in [0.,.02,.04] for z in [.7,1.3]]
        bounds=([0.,5.,.3,.8,-.3],[.12,150.,4.,1.2,.3])
    best=None
    for start in starts:
        fit=least_squares(lambda x:(simulate(ct,command,t,unpack(x))-measured)[train],
                          start,bounds=bounds,max_nfev=180,diff_step=1e-3)
        if best is None or fit.cost<best.cost:
            best=fit
    params=unpack(best.x)
    predicted=simulate(ct,command,t,params)
    lag=params['delay_s']+(params['tau_s'] if kind=='first' else 2*params['zeta']/params['wn_rad_s'])
    return params,predicted,lag


def identify(folder):
    d=json.loads((folder/'analysis.json').read_text());p=d['plots']
    j=int(d['metadata']['config'].get('joint',6))-1
    ct=np.asarray(p['command_time']);t=np.asarray(p['time'])
    command=np.asarray(p['command_delta'])[:,j]
    measured=np.asarray(p['actual_delta'])[:,j]
    train=(t>=0)&(t<18)
    valid=(t>=18)&(t<=ct[-1])
    dynamic=np.zeros(len(t),dtype=bool)
    for s in d['steps']:
        if s['joint']==j+1:
            dynamic|=(t>=s['start_s'])&(t<s['start_s']+.7)
    train &= dynamic
    valid &= dynamic
    models=[]
    for kind in ['first','second']:
        params,y,lag=fit_model(ct,command,t,measured,train,kind)
        models.append(dict(params=params,parameter_output_units='mrad (bias only)',
            train_rmse_mrad=float(np.sqrt(np.mean((y[train]-measured[train])**2))),
            validation_rmse_mrad=float(np.sqrt(np.mean((y[valid]-measured[valid])**2))),
            validation_max_error_mrad=float(np.max(np.abs(y[valid]-measured[valid]))),
            steady_ramp_lag_s=lag,predicted_mrad=y.tolist()))
    result=dict(joint=j+1,train='2 and 4 mrad transitions; first 0.7 s after each step',
        validation='8 mrad transitions, never used for fitting; first 0.7 s after each step',
        limitations=['Single joint and posture, small free-space movement.',
          'Bias is in mrad; convert to radians when used in a joint-space simulator.',
          'Fitted delay includes clock-alignment and transport effects; not a pure network measurement.',
          'Not identified: torque limits, contact response, friction, payload effects or saturation.',
          'Do not apply this filter plus an already-laggy actuator and double-count the response.'],
        time_s=t.tolist(),command_time_s=ct.tolist(),command_mrad=command.tolist(),
        measured_mrad=measured.tolist(),models=models)
    (folder/'identified_models.json').write_text(json.dumps(result,indent=2))
    for model in models:
        print(json.dumps({k:v for k,v in model.items() if k!='predicted_mrad'}))
    return result


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('folder',type=Path)
    identify(parser.parse_args().folder)