
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from welib.fast.campbell import campbellModeStyles
from welib.system.statespace import StateMatrix
from welib.system.eva import eig, eigA

from welib.system.wtmodels.model15M import *

scriptDir=os.path.dirname(__file__)
outDir=os.path.join(scriptDir,'_outputs')
os.makedirs(outDir, exist_ok=True)


np.set_printoptions(linewidth=300, precision =5)


def main(test=False):
    vOmega=np.arange(0,1.55,0.05) #omegas up to rate omega Rotor Speed
    # vOmega=np.arange(0,10.55,0.05) #omegas up to rate omega Rotor Speed
    p=parameters() # Turbine parameters
    lambda1 = 1.8751 # NOTE: constant for analytical bending mode 1
    c1      = 0.7341 # NOTE: constant for analytical bending mode 1, see bendingMode!

    df=pd.DataFrame()
    df['Omega_[rad/s]'] = vOmega
    df['f_flap_[Hz]']   = vOmega*0
    df['f_edge_[Hz]']   = vOmega*0
    df['f_tors_[Hz]']   = vOmega*0

    for i,Omega in enumerate(vOmega):
        Mb,Kb,Gb,Mt,Gt,Kt = diagonalStructuralMatrices(p,Omega,lambda1,c1)
        A=StateMatrix(M=Mb,K=Kb)
        Q,freq =eig(K=Kb,M=Mb, freq_out=True, sort=True)
        freq2, zeta, Q2, freq_0  = eigA(A, nq=Mb.shape[0])
        df.loc[i,'f_flap_[Hz]'] = freq[0]
        df.loc[i,'f_edge_[Hz]'] = freq[1]
        df.loc[i,'f_tors_[Hz]'] = freq[2]

    if not test:
        df.to_csv(os.path.join(outDir,'A_freq.csv'), sep='\t', index=False)


    # --- Plot 
    fig,ax = plt.subplots(1, 1, sharey=False, figsize=(6.4,4.8)) # (6.4,4.8)
    fig.subplots_adjust(left=0.12, right=0.95, top=0.75, bottom=0.11, hspace=0.20, wspace=0.20)
    lbl='1st blade flap' # TODO Mode 1
    c, ls, ms, mk = campbellModeStyles(0, lbl)
    ax.plot(vOmega, df['f_flap_[Hz]'], ls, marker=mk, color=c, label=lbl)

    lbl='1st blade edge' # TODO Mode 2
    c, ls, ms, mk = campbellModeStyles(1, lbl)
    ax.plot(vOmega, df['f_edge_[Hz]'], ls, marker=mk, color=c, label=lbl)

    lbl='1st blade torsion' # TODO TODO
    c, ls, ms, mk = campbellModeStyles(2, lbl)
    ax.plot(vOmega, df['f_tors_[Hz]'], ls, marker=mk, color=c, label=lbl)

    ax.set_xlabel(r'$\Omega$ [rad/s]')
    ax.set_ylabel(r'Frequency [Hz]')
    ax.legend()
    if not test:
        fig.savefig(os.path.join(outDir,'ExA1_ModeFreqs.png'))


    # --------------------------------------------------------------------------------}
    # ---  
    # --------------------------------------------------------------------------------{
    x=np.linspace(0,50,51)
    nDOF= 3
    phi_e = bendingMode(x,1)  
    phi_t = torsionalMode(x,1)
    phi_f = bendingMode(x,1)  
    Phi = np.column_stack((phi_f, phi_e, phi_t))
    # for j in np.arange(nDOF):
    #     Phi[:,j] /= Phi[-1,j]

    # print(Phi)

    # --- Eigen vector 
    def getShape(Omega):
        Mb,Kb,Gb,Mt,Gt,Kt = diagonalStructuralMatrices(p,Omega,lambda1,c1)
        A = StateMatrix(M=Mb,K=Kb)
        freq2, zeta, Q, freq_0  = eigA(A, nq=Mb.shape[0], normQ='byMax')

        Q0,freq0 =eig(K=Kb,M=Mb, freq_out=True, sort=True, normQ='byMax')
        phase = np.angle(Q)
        ampl  = np.abs(Q)
        #print(phase)
        #print(ampl)
        Modes = np.zeros((Q.shape[1], Phi.shape[0], Phi.shape[1]))
        for j in np.arange(Q.shape[1]):
            q_j = Q[:,j]
            #a_j = np.abs(q_j) # TODO figure out how to incorpoorate sign with phase
            phi_j = np.angle(q_j) 
            a_j = Q0[:,j] 
            Mode = np.multiply(Phi, a_j)
            Mode_end=Mode[-1,:]
            scale = np.max(Mode_end)
            Modes[j][:,:] = Mode/scale
        return Modes, phase

    Modes0,phase0 = getShape(Omega=0)   
    ModesR,phaseR = getShape(Omega=1.5)

    for j in np.arange(nDOF):
        Mode0=Modes0[j]
        ModeR=ModesR[j]
        fig,ax = plt.subplots(1, 1, sharey=False, figsize=(6.4,4.8)) # (6.4,4.8)
        fig.subplots_adjust(left=0.12, right=0.95, top=0.95, bottom=0.11, hspace=0.20, wspace=0.20)
        ax.plot(x, Mode0[:,0], label='flap')
        ax.plot(x, Mode0[:,1], label='edge')
        ax.plot(x, Mode0[:,2], label='torsion')
        ax.plot(x, ModeR[:,0], 'kx')
        ax.plot(x, ModeR[:,1], 'kx')
        ax.plot(x, ModeR[:,2], 'kx')
        ax.set_xlabel('Blade span [m]')
        ax.set_ylabel('Mode content')
        ax.set_title('Mode {}'.format(j+1))
        ax.legend()
        if not test:
            fig.savefig(os.path.join(outDir,'ExA1_Mode{}.png'.format(j+1)))
    return df, ModesR
# print(Omega)
# print('Mb\n',Mb)
# print('Kb\n',Kb)
# print('Mt\n',Mt)
# print('Gt\n',Gt)
# print('Kt\n',Kt)
# print('Q\n',Q)
# print('freq\n',freq)
# print('Q\n',np.imag(Q2))
# print('freq_d\n',freq_d)
# fig,ax = plt.subplots(1, 1, sharey=False, figsize=(6.4,4.8)) # (6.4,4.8)
# fig.subplots_adjust(left=0.12, right=0.95, top=0.95, bottom=0.11, hspace=0.20, wspace=0.20)
# ax.plot(x, phi_e    , label='')
# ax.set_xlabel('')
# ax.set_ylabel('')
# ax.legend()

if __name__ == '__main__':
    df, ModesR = main()
    print(df)
    plt.show()


if __name__ == '__test__':
    df, ModesR = main(test=True)
    np.testing.assert_almost_equal(df['f_flap_[Hz]'].values[:3], [1.386216,1.386236,1.386296],5) 
    np.testing.assert_almost_equal(df['f_tors_[Hz]'].values[:3], [4.139987,4.140006,4.140064],5) 
    np.testing.assert_almost_equal(ModesR[2][-4:,0], [0.917416, 0.944941, 0.97247 , 1.], 6)
    np.testing.assert_almost_equal(ModesR[2][1:4,0], [0.0006967, 0.0027612, 0.0061546], 6)
    np.testing.assert_almost_equal(ModesR[2][1:4,2], [0.0167179, 0.0334194, 0.0500879], 5)
