%run matt_startup
%run -i matt_utils

button_qtconsole()

#import other needed modules in all used engines
#with dview.sync_imports():
#   import os

"""
Explanation: Parameter Estimation of RIG Roll Experiments
Setup and descriptions

Without ACM model
Turn on wind tunnel
Only 1DoF for RIG roll movement
Use small-amplitude aileron command of CMP as inputs (in degrees)
$$U = \delta_{a,cmp}(t)$$
Consider RIG roll angle and its derivative as States (in radians)
$$X = \begin{pmatrix} \phi_{rig} \ \dot{\phi}_{rig} \end{pmatrix}$$
Observe RIG roll angle and its derivative as Outputs (in degrees)
$$Z = \begin{pmatrix} \phi_{rig} \ \dot{\phi}_{rig} \end{pmatrix}$$
Use output error method based on most-likehood(ML) to estimate
$$ \theta = \begin{pmatrix} C_{l,\delta_a,cmp} \ C_{lp,cmp}  \end{pmatrix} $$

Startup computation engines
End of explanation
"""


filename = 'FIWT_Exp015_20150601145005.dat.npz'

def loadData():
    # Read and parse raw data
    global exp_data
    exp_data = np.load(filename)

    # Select colums
    global T_cmp, da_cmp
    T_cmp = exp_data['data33'][:,0]
    da_cmp = np.average(exp_data['data33'][:,3:11:2], axis=1)

    global T_rig, phi_rig
    T_rig = exp_data['data44'][:,0]
    phi_rig = exp_data['data44'][:,2]

loadData()

text_loadData()

"""
Explanation: Data preparation
Load raw data
End of explanation
"""


def checkInputOutputData():

    #check inputs/outputs
    fig, ax = plt.subplots(2,1,True)
    ax[0].plot(T_cmp,da_cmp,'r',picker=1)
    ax[1].plot(T_rig,phi_rig, 'b', picker=2)
    ax[0].set_ylabel('$\delta \/ / \/ ^o$')
    ax[1].set_ylabel('$\phi \/ / \/ ^o/s$')
    ax[1].set_xlabel('$T \/ / \/ s$', picker=True)
    ax[0].set_title('Output', picker=True)
    

    fig.canvas.mpl_connect('pick_event', onPickTime)
    fig.show()
    display(fig)

button_CheckData()

"""
Explanation: Check time sequence and inputs/outputs
Click 'Check data' button to show the raw data. 
Click on curves to select time point and push into queue; click 'T/s' text to pop up last point in the queue; and click 'Output' text to print time sequence table.
End of explanation
"""



# Pick up focused time ranges
time_marks = [[1501.28, 1505.50, "doublet u1"],
              [1507.40, 1511.80, "doublet u2"],
              [1513.55, 1517.87, "doublet u3"],
              [1519.70, 1523.50, "doublet u4"],
              [1537.60, 1541.64, "doublet d1"],
              [1543.76, 1547.90, "doublet d2"],
              [1549.91, 1554.20, "doublet d3"],
              [1555.86, 1560.00, "doublet d4"],
              [1609.30, 1615.49, "3-2-1-1 u1"], 
              [1617.89, 1624.25, "3-2-1-1 u2"],
              [1626.49, 1633.45, "3-2-1-1 u3"],
              [1634.99, 1642.38, "3-2-1-1 u4"],
              [1651.40, 1657.50, "3-2-1-1 d1"],
              [1659.90, 1666.68, "3-2-1-1 d2"],
              [1668.50, 1674.69, "3-2-1-1 d3"],
              [1677.00, 1683.88, "3-2-1-1 d4"],
              [1748.59, 1809.05, "linear sweep u1"], 
              [1825.89, 1885.96, "linear sweep d1"], 
              [1905.86, 1965.17, "exp sweep u1"], 
             ]


# Decide DT,U,Z and their processing method
DT=0.01
process_set = {
            'U':[(T_cmp, da_cmp,0),],
            'Z':[(T_rig, phi_rig,3),],
            'cutoff_freq': 10 #Hz
           }

U_names = ['$\delta_{a,cmp} \, / \, ^o$',]
Y_names = Z_names = ['$\phi_{a,rig} \, / \, ^o$',
                    '$\dot{\phi}_{a,rig} \, / \, ^o/s$',]

display_data_prepare()

"""
Explanation: Input $\delta_T$ and focused time ranges
For each section,
* Select time range and shift it to start from zero;
* Resample Time, Inputs, Outputs in unique $\delta_T$;
* Smooth Input/Observe data if flag bit0 is set;
* Take derivatives of observe data if flag bit1 is set.
End of explanation
"""


resample(True);

"""
Explanation: Resample and filter data in sections
End of explanation
"""


%%px --local
#update common const parameters in all engines

#problem size
Nx = 2
Nu = 1
Ny = 2
Npar = 9

#reference
S_c = 0.1254   #S_c(m2) 
b_c = 0.7      #b_c(m)
g = 9.81       #g(m/s2)
V = 30         #V(m/s)

#other parameters
v_th = 0.5/57.3  #v_th(rad/s)
v_th2 = 0.5/57.3  #v_th(rad/s)

#for short
qbarSb = 0.5*1.225*V*V*S_c*b_c
b2v = b_c/(2*V)

def x0(Z,T,U,params):
    return Z[0,:]/57.3

def xdot(X,t,U,params):
    Cla_cmp = params[0]
    Clp_cmp = params[1]
    Ixx     = params[2]
    F_c     = params[3]
    f       = params[4]
    m_T     = params[5]
    l_z_T   = params[6]
    kBrk    = params[7]
    phi0    = params[8]

    phi     = X[0]
    phi_dot = X[1]
    
    idx = int(t/DT)
    da_cmp = U[idx,0]*0.01745329
    
    moments = -(m_T*l_z_T)*g*math.sin(phi-phi0)

    abs_phi_dot = abs(phi_dot)
    F = f*phi_dot
    if abs_phi_dot > v_th+v_th2:
        F += math.copysign(F_c, phi_dot)
    elif abs_phi_dot > v_th:
        F += math.copysign(F_c*(kBrk-(kBrk-1)*(abs_phi_dot-v_th)/v_th2), phi_dot)
    else:
        F += phi_dot/v_th*(F_c*kBrk)
    moments -= F
    
    moments += qbarSb*(Cla_cmp*da_cmp + Clp_cmp*phi_dot*b2v)
    
    phi_dot2 = moments/Ixx
    return [phi_dot, phi_dot2]

def obs(X,T,U,params):
    return X*57.3



display(HTML('<b>Constant Parameters</b>'))
table = ListTable()
table.append(['Name','Value','unit'])
table.append(['$S_c$',S_c,r'$m^2$'])
table.append(['$b_c$',b_c,'$m$'])
table.append(['$g$',g,r'$ms^{-2}$'])
table.append(['$V$',V,r'$ms^{-1}$'])
display(table)


"""
Explanation: Define dynamic model to be estimated
$$\left{\begin{matrix}
\ddot{\phi}{rig} = \frac{M{x,rig}}{I_{xx,rig}} \
M_{x,rig} = M_{x,a} + M_{x,f} + M_{x,cg} \
M_{x,a} = \frac{1}{2} \rho V^2S_cb_c \left ( C_{la,cmp}\delta_{a,cmp} + C{lp,cmp} \frac{b_c}{2V} \dot{\phi}{rig} \right ) \
M{x,f} = -F_c \, sign(\dot{\phi}{rig}) - f\dot{\phi}{rig} \
M_{x,cg} = -m_T g  l_{zT} \sin \left ( \phi - \phi_0 \right ) \
\end{matrix}\right.$$
End of explanation
"""


#initial guess
param0 = [
    -0.3,      #Cla_cmp(1/rad)
    -0.5,      #Clp_cmp
    0.199141909329,    #Ixx(kg*m2)
    0.0580817418532,      #F_c(N*m)
    0.0407466009837,      #f(N*m/(rad/s))
    7.5588,     #m_T(kg)
    0.0444,    #l_z_T(m)
    1.01,     #kBrk
       0,      #phi0(rad)
]

param_name = ['$Cla_{cmp}$',
              '$Clp_{cmp}$',
              '$I_{xx,rig}$', 
              '$F_c$',
              '$f$',
              '$m_T$',
              '$l_{zT}$',
              '$k_{Brk}$',
              '$phi_0$']
param_unit = ['$rad^{-1}$',
              '$1$',
              '$kg\,m^2$',
              '$Nm$',
              r'$\frac{Nm}{rad/s}$',
              'kg',
              'm',
              '1',
              '$rad$']
NparID = 4
opt_idx = [0,1,2,8]
opt_param0 = [param0[i] for i in opt_idx]
par_del = [0.3*1e-3,  0.3*1e-3,   0.2*1e-3,   0.0174]
bounds =  [(-1,-1e-6),(-2,-1e-6), (1e-6,0.5), (-0.1,0.1)]
'''
NparID = 3
opt_idx = [0,1,7]
opt_param0 = [param0[i] for i in opt_idx]
par_del = [0.3*1e-3,  0.3*1e-3,   0.0174]
bounds =  [(-1,-1e-6),(-2,-1e-6), (-0.1,0.1)]
'''


display_default_params()

#select sections for training
section_idx = range(8)
display_data_for_train()

#push parameters to engines
push_opt_param()


# select 2 section from training data
idx = random.sample(section_idx, 2)

interact_guess();

"""
Explanation: Initial guess

Input default values and ranges for parameters
Select sections for trainning
Adjust parameters based on simulation results
Decide start values of parameters for optimization
End of explanation
"""


display_preopt_params()

if True:
    InfoMat = None
    method = 'trust-ncg'
    def hessian(opt_params, index):
        global InfoMat
        return InfoMat
    dview['enable_infomat']=True
    options={'gtol':1}
    opt_bounds = None
else:
    method = 'L-BFGS-B'
    hessian = None
    dview['enable_infomat']=False
    options={'ftol':1e-2,'maxfun':10}
    opt_bounds = bounds

cnt = 0
tmp_rslt = None
T0 = time.time()
print('#cnt,     Time,      |R|')

%time res =  sp.optimize.minimize(fun=costfunc, x0=opt_param0, \
        args=(opt_idx,), method=method, jac=True, hess=hessian, \
        bounds=opt_bounds, options=options)


"""
Explanation: Optimize using ML
End of explanation
"""


display_opt_params()

# show result
idx = random.sample(range(8), 2) \
    + random.sample(range(8,16), 2) \
    + random.sample(range(16,19), 2)
display_data_for_test();

update_guess();

toggle_inputs()
button_qtconsole()

"""
Explanation: Show and test results
End of explanation
"""


