""" 
Manages packages generated by YAMS

"""
import os
import numpy as np
import pandas as pd
import importlib

def loadPackage(modelName=None, packageDir='', packagePath=None):
    """ """
    if modelName is not None:
        packagePath = os.path.join(packageDir, modelName)+'.py'
    else:
        modelName = os.path.basename(packagePath)
    spec = importlib.util.spec_from_file_location(modelName, packagePath)
    if spec is None:
        raise Exception('Package not found: ',packagePath)
    print('>>> Loading package', packagePath)
    pkg = spec.loader.load_module()
    return pkg, packagePath



# --------------------------------------------------------------------------------}
# --- Nonlinear simulation 
# --------------------------------------------------------------------------------{
def simulate(pkg, time, q0, qd0=None, p=None, u=None, acc=False, forcing=False, **kwargs):
    """ 
    Perform nonlinear simulation based on a `model` (python package) generated by yams_sympy
     - time: time vector
     - q0:  initial conditions for positions
     - qd0: initial conditions for speeds
     - p: parameters
     - u: inputs
     - acc: compute acceleration afterwards and store in datadrame
     - forcing: compute forcing afterwards and store in dataframe
     - kwargs (see MechSystem.toDataFrame)
    """
    from welib.system.mech_system import MechSystem

    if qd0 is None:
        qd0=q0*0
    if u is None:
        u=dict()
    for key in pkg.info()['su']:
        if key in u.keys():
            print('[INFO] packman.py  u[{}] set by user'.format(key))
        else:
            # TODO get it from ref simulation
            print('[WARN] packman.py: u[{}] set to 0'.format(key))
            u[key]= lambda t,q=None,qd=None: 0

    # --- Test to see if model works
    MM0      = pkg.mass_matrix(q0,p)
    forcing0 = pkg.forcing(0,q0,qd0,p,u)

    # --- integrate nonlinear system
    fM = lambda x: pkg.mass_matrix(x, p)
    fF = lambda t,x,xd: pkg.forcing(t, x, xd, p=p, u=u)
    sysNL = MechSystem(fM, F=fF, x0=q0, xdot0=qd0 )
    resNL = sysNL.integrate(time, method='RK45')

    # Storing
    sysNL._M0       = MM0
    sysNL._forcing0 = forcing0.ravel()

    # Create a Mechanical System
    sysNL = MechSystem(fM, F=fF, x0=q0, xdot0=qd0 )
    # Integrate in time
    resNL, _ = sysNL.integrate(time, method='RK45')
    # Store in a nice dataframe, potentially with acceleration and forcing
    calc=''
    if acc:
        calc+='xdd,'
    if forcing:
        calc+='f,'
    dfNL = sysNL.res2DataFrame(resNL, calc=calc, **kwargs)

    # TODO point accelerations, additional outputs etc.

    return resNL, sysNL, dfNL


# --------------------------------------------------------------------------------}
# --- Linear model 
# --------------------------------------------------------------------------------{
def linearModel(pkg, p, dq0=None, dqd0=None, time=None, uop=None, qop=None, qdop=None, du=None, MCKextra=None, MCKu=None, noBlin=False, sX=None, sXd=None):
    """ Perform linear simulation based on a model (python package) generated by yams_sympy """
    # TODO TODO TODO MOVE ME TO packman
    from welib.system.mech_system import MechSystem
    from welib.tools.signal_analysis import interpArray

    info = pkg.info()

    # --- Perturbations about operating point
    if dq0 is None:
        dq0=np.zeros(info['nq'])
    if dqd0 is None:
        dqd0=np.zeros(info['nq'])
    # --- Operating point
    if qop is None:
        qop = dq0*0 
    if qdop is None:
        qdop= dqd0*0 
    if uop is None:
        print('[WARN] py_lin: Setting all input operating points to 0')
        uop=dict() # Inputs at operating points
        for key in info['su']:
            uop[key]= 0 

    # --- Initial conditions (with op)
    q0  = dq0 + qop
    qd0 = dqd0 + qdop
#     dq0  = self.q0  - qop
#     dqd0 = self.qd0 - qdop
    print('q0  :',q0)
    print('qd0 :',qd0)
    print('qop :',qop)
    print('qdop:',qdop)
    print('uop: ',uop)
    #print('dq0 :',dq0)
    #print('dqd0:',dqd0)

    # --- Evaluate linear structural model at operating point
    M_lin   = pkg.M_lin(qop,p)
    C_lin   = pkg.C_lin(qop,qdop,p,uop)
    K_lin   = pkg.K_lin(qop,qdop,p,uop) 
    B_lin   = pkg.B_lin(qop,qdop,p,uop)

    # --- Integrate linear system
    if MCKextra is not None:
        Me, Ce, Ke = MCKextra
        M_lin += Me
        C_lin += Ce
        K_lin += Ke
    print('B_lin\n',B_lin)
    if noBlin:
        n,m=B_lin.shape
        B_lin = np.eye(max(n,m))[:n,:m]
        print('B_lin\n',B_lin)

    if MCKu is not None:
        Mu, Cu, Ku = MCKu
        M_lin += B_lin.dot(Mu)
        C_lin += B_lin.dot(Cu)
        K_lin += B_lin.dot(Ku)
        Z = np.zeros(B_lin.shape[0])


    # --- Time inputs
    Z = np.zeros(B_lin.shape[0])
    fF  = lambda t,x=None,xd=None: Z
    if time is not None:
        if MCKu is None:
            if len(time)>0:
                if du is None:
                    print('[WARN] linearModel: Setting all input time series to 0')
                    nu = len(info['su'])
                    du = np.zeros((nu, len(time)))
                fF  = lambda t,x=None,xd=None: B_lin.dot( interpArray(t, time, du) )
    #fdu = lambda t,x=None,xd=None: interpArray(t, time, du)

    forcing0=fF(0)

    # --- Setup Mech system (for time integration)
    sysLI = MechSystem(M=M_lin, K=K_lin, C=C_lin, F=fF, x0=dq0, xdot0=dqd0, sX=sX, sXd=sXd)
    #print(sysLI)
    #print('nu',nu)
    #print('du',du)
    #print('fF',fF(0))
    #print('fF',fF(0,dq0))
    # --- Store
    sysLI._B    = B_lin
    sysLI._qop  = qop
    sysLI._qdop = qdop
    sysLI._uop  = uop
    sysLI._forcing0  = forcing0.ravel()

    return sysLI



# --------------------------------------------------------------------------------}
# ---  
# --------------------------------------------------------------------------------{
def stateAcc(q,qd,F,p,u,pkg):
    M=pkg.mass_matrix(q,p)
    #K=pkg.K_lin(q,qd,p)
    #C=pkg.C_lin(q,qd,p)
    F=pkg.forcing(0,q,qd,p,u).flatten()
    Minv= np.linalg.inv(M)
    #qdd= -Minv.dot(K).dot(q) - Minv.dot(C).dot(qd)
    qdd= Minv.dot(F)
    return qdd

def AccelerationP_IMU(q, qd, p, u, pkg):
    """ Acceleration of IMU as function of q and qd only"""
    qdd = stateAcc(q,qd,None,p,u,pkg)
    acc = pkg.AccelerationP_IMU(q, qd, qdd, p).flatten()
    return acc


# --------------------------------------------------------------------------------}
# --- Compute Jacobian of IMU acc wrt states 
# --------------------------------------------------------------------------------{
def IMUjacobian(pkg, q0, qd0, p, method='finiteDifferences', uop=None, u=None, dq=None, dqd=None, sDOFs=None):
    """ 
    Compute Jacobian of IMU (acceleration at one point) with respect to states
    based on the functions provided in the package

    Example of calculation for dq and dqd
             dq  = ((np.max(dfFS[sq]) -np.min(dfFS[sq]))/100).values
             dqd = ((np.max(dfFS[sqd])-np.min(dfFS[sqd]))/100).values
    """
    if qd0 is None:
        qd0=np.asarray(q0)*0

    fh = lambda q,qd : AccelerationP_IMU(q,qd,p,u,pkg)
    f0 = fh(q0, qd0) #, p, u, pkg)
    if method=='finiteDifferences':
        # --- Method 1 Finite differences
        from welib.system.linearization import numerical_jacobian
        qdd0 = q0*0
        if dq is None or dqd is None:
            raise Exception('dq and dqd need to be provided when using finite differences')
        Kacc = numerical_jacobian(fh, (q0,qd0), 0, dq )
        Cacc = numerical_jacobian(fh, (q0,qd0), 1, dqd)
    elif method=='packageJacobians':
        # --- Method 2 pkg
        Ma,Ca,Ka = pkg.AccLinP_IMU(q=q0, qd=qd0, p=p)
        M    = pkg.M_lin(q0,p)
        Minv = np.linalg.inv(M)
        K    = pkg.K_lin(q0,qd0,p,uop)
        C    = pkg.C_lin(q0,qd0,p,uop)
        Kacc = Ka - Ma.dot(Minv).dot(K)
        Cacc = Ca - Ma.dot(Minv).dot(C)
    else:
        raise NotImplementedError()

    if sDOFs is not None:
        Kacc = pd.DataFrame(data=Kacc, columns=sDOFs, index=['NcIMUTAxs_[m/s^2]','NcIMUTAys_[m/s^2]','NcIMUTAzs_[m/s^2]']) 
        sDOFsd = ['d'+s for s in sDOFs]
        Cacc = pd.DataFrame(data=Cacc, columns=sDOFsd, index=['NcIMUTAxs_[m/s^2]','NcIMUTAys_[m/s^2]','NcIMUTAzs_[m/s^2]']) 

    return Kacc, Cacc, f0

