# -*- coding: utf-8 -*-
"""
Created on Tue Sep 12 18:47:50 2017

@author: adelpret
"""
import numpy as np
from scipy import ndimage
import matplotlib.pyplot as plt
from identification_utils import solve1stOrderLeastSquare
from dynamic_graph.sot.torque_control.hrp2.control_manager_conf import IN_OUT_GAIN


def identify_motor_static(enc, dq, ctrl, current, tau, JOINT_ID, JOINT_NAME, ZERO_VELOCITY_THRESHOLD, 
                          ZERO_VELOCITY_THRESHOLD_SMALL, SHOW_THRESHOLD_EFFECT):
    # remove high velocity
    maskConstAng = (abs (dq)<ZERO_VELOCITY_THRESHOLD)
    # erode to get only steady phases where velocity is small 
    maskConstAng=ndimage.morphology.binary_erosion(maskConstAng,None,100)
    maskPosVel=(dq> ZERO_VELOCITY_THRESHOLD_SMALL)
    maskNegVel=(dq<-ZERO_VELOCITY_THRESHOLD_SMALL)
    maskConstPosAng=np.logical_and( maskConstAng ,maskPosVel )
    maskConstNegAng=np.logical_and( maskConstAng ,maskNegVel ) 
    if SHOW_THRESHOLD_EFFECT :
        plt.figure()
        plt.plot(enc, label='q')
        q_const=enc.copy()
        q_const[np.logical_not(maskConstAng)]=np.nan
        plt.plot(q_const, label='q_const')
        plt.legend()
        
    # identify current sensor gain
    x = current[maskConstAng]
    y = ctrl[maskConstAng]/IN_OUT_GAIN
    maskPosErr = np.logical_and(y - x > 0.0, np.abs(x)>0.5)
    maskNegErr = np.logical_and(y - x < 0.0, np.abs(x)>0.5)
    print "Number of samples with constant angle:", x.shape[0]
    print "Number of samples with constant angle and pos vel:", x[maskPosErr].shape[0]
    print "Number of samples with constant angle and neg vel:", x[maskNegErr].shape[0]
    if(x[maskPosErr].shape[0]<10):
        (Ks,DZ)=solve1stOrderLeastSquare(x[maskNegErr], y[maskNegErr])
    elif(x[maskNegErr].shape[0]<10):
        (Ks,DZ)=solve1stOrderLeastSquare(x[maskPosErr], y[maskPosErr])
    else:
        (Ksn,DZn)=solve1stOrderLeastSquare(x[maskNegErr], y[maskNegErr])
        (Ksp,DZp)=solve1stOrderLeastSquare(x[maskPosErr], y[maskPosErr])    
        Ks = 0.5*(Ksp+Ksn);
        Ks = min([Ksp, Ksn]);
        DZ = 0.5*(DZp-DZn);    
        print "Current sensor gains = ", Ksp, Ksn;
        print "Deadzones            = ", DZp, -DZn;
        
        x_neg = x[maskNegErr]
        y_neg = y[maskNegErr]
        plt.figure()    
        plt.plot(x_neg, y_neg,'.' ,lw=3,markersize=1,c='0.5');  
        plt.plot([min(x_neg),max(x_neg)],[Ksn*min(x_neg)+DZn ,Ksn*max(x_neg)+DZn],'g:',lw=3)
        plt.ylabel(r'$i(t)$'); plt.xlabel(r'$u(t)$')
        plt.title('Negative current errors - Joint '+JOINT_NAME)
        
        x_pos = x[maskPosErr]
        y_pos = y[maskPosErr]
        plt.figure()    
        plt.plot(x_pos, y_pos,'.' ,lw=3,markersize=1,c='0.5');  
        plt.plot([min(x_pos),max(x_pos)],[Ksp*min(x_pos)+DZp ,Ksp*max(x_pos)+DZp],'g:',lw=3)
        plt.ylabel(r'$i(t)$'); plt.xlabel(r'$u(t)$')
        plt.title('Positive current errors - Joint '+JOINT_NAME)
        plt.show()
    
    if(Ks<0.0):
        print "ERROR: estimated Ks is negative! Setting it to 1"
        Ks = 1.0;
    
    # plot dead zone effect ********************************************
    plt.figure()
    plt.plot(Ks*current, label='current')
    plt.plot(ctrl/IN_OUT_GAIN, label='control')
    plt.legend()
    
    plt.figure()
    y = Ks*current[maskConstAng]
    x = ctrl[maskConstAng]/IN_OUT_GAIN - Ks*current[maskConstAng]
    plt.ylabel(r'$i(t)$')
    plt.xlabel(r'$ctrl(t)-i(t)$')
    plt.plot(x,y,'.' ,lw=3,markersize=1,c='0.5');  
    plt.plot(x[maskPosErr],y[maskPosErr],'rx',lw=3,markersize=1, label='pos err'); 
    plt.plot(x[maskNegErr],y[maskNegErr],'bx',lw=3,markersize=1, label='neg err'); 
    plt.legend()
    
    plt.figure()
    y = ctrl[maskConstAng]/IN_OUT_GAIN
    x = ctrl[maskConstAng]/IN_OUT_GAIN - Ks*current[maskConstAng]
    plt.ylabel(r'$ctrl(t)$')
    plt.xlabel(r'$ctrl(t)-i(t)$')    
    plt.plot(x,y,'.' ,lw=3,markersize=1,c='0.5');  
    plt.plot(x[maskPosErr],y[maskPosErr],'rx',lw=3,markersize=1, label='pos err'); 
    plt.plot(x[maskNegErr],y[maskNegErr],'bx',lw=3,markersize=1, label='neg err'); 
    plt.legend()
    
    plt.figure()
    y = ctrl/IN_OUT_GAIN
    x = Ks*current
    plt.ylabel(r'$ctrl(t)$')
    plt.xlabel(r'$i(t)$')    
    plt.plot(x,y,'.' ,lw=3,markersize=1,c='0.5');  
    plt.plot([-3,3],[-3,3]);  
    
    plt.show()
#    y = a. x   +  b
#    i = Kt.tau + Kf
    
    # Identification ***************************************************
    y = current #*Ks
    x = tau
    (Ktp,Kfp)=solve1stOrderLeastSquare(x[maskConstPosAng],y[maskConstPosAng])
    (Ktn,b)=solve1stOrderLeastSquare(x[maskConstNegAng],y[maskConstNegAng])
    Kfn=-b
    
    # Plot *************************************************************
    plt.figure()    
    plt.axhline(0, color='black',lw=1)
    plt.axvline(0, color='black',lw=1)
    plt.plot(x     ,y     ,'.' ,lw=3,markersize=1,c='0.5');  
    plt.plot(x[maskConstPosAng],y[maskConstPosAng],'rx',lw=3,markersize=1); 
    plt.plot(x[maskConstNegAng],y[maskConstNegAng],'bx',lw=3,markersize=1); 
    #plot identified lin model
    plt.plot([min(x),max(x)],[Ktp*min(x)+Kfp ,Ktp*max(x)+Kfp],'g:',lw=3)
    plt.plot([min(x),max(x)],[Ktn*min(x)-Kfn ,Ktn*max(x)-Kfn],'g:',lw=3)
    plt.ylabel(r'$i(t)$')
    plt.xlabel(r'$\tau(t)$')
    plt.title('Static experiment - Joint '+JOINT_NAME)

    print "cur_sens_gain[%d] = %f" % (JOINT_ID, Ks);
    print 'deadzone[%d]      = %f' % (JOINT_ID, DZ);    
    print 'Kt_p[%d]          = %f' % (JOINT_ID,Ktp);
    print 'Kt_n[%d]          = %f' % (JOINT_ID,Ktn);
    print 'Kf_p[%d]          = %f' % (JOINT_ID,Kfp);
    print 'Kf_n[%d]          = %f' % (JOINT_ID,Kfn);
    
    print 'Kt_m[%d]          = %f' % (JOINT_ID,(Ktp+Ktn)/2.0);
    print 'Kf_m[%d]          = %f' % (JOINT_ID,(Kfp+Kfn)/2.0);
    
    return (Ktp, Ktn, Ks, DZ);