import numpy as np
import matplotlib.pyplot as plt
import scipy.special as spe
from scipy import linalg
import pseudo_spectral_v2_revised as ps2
import pseudo_spectral_v3_revised as ps3
import pseudo_spectral_v4_revised as ps4
import test2
import pandas as pd
import acc
import initial_input as ip


Lm = np.pi
num_of_VA = 1
num_of_Lm = 10
# Lm_list = np.linspace(np.pi/2+0.2,2*np.pi,num_of_Lm)
Lm_list = [Lm]
y_collo = spe.roots_hermite(ip.n_order)[0]


# V_A_list = np.concatenate((np.linspace(0,0.08,9),np.linspace(0.081,0.09,10),np.linspace(0.1,0.3,21)))
# V_A_list = np.concatenate((np.linspace(0,0.03,4),np.linspace(0.031,0.04,10),np.linspace(0.05,0.2,16)))
# V_A_list = np.linspace(0,0.3,num_of_VA)
V_A_list = [0.20]
tmp_max = np.zeros(num_of_VA)
width_Lm = np.zeros((num_of_Lm,num_of_VA))
U0 = ip.Q0


j = 0
for Lm in Lm_list:
    tmp_max = np.zeros(num_of_VA)
    kelvin_i = np.zeros(num_of_VA)
    doppler_shift = np.zeros((num_of_VA,3,2)) # the argument and modulus of coefficient of kelvin mode(-1) and Rossby 1 ,3

    # rossby1_10 = np.zeros((6,num_of_VA)) # eigenfrequency of Rossby mode (j = 1,3,5,7,9)
    for i in range(num_of_VA):
        print('.',end = ' ')
        ##------------------------------ get the data -----------------------##
        V_A = V_A_list[i]
        # Lm = Lm_list[i]
        
        phi40_contra = test2.phi_collo_array(y_collo)
        # H1,f1 = ps2.assign(ip.n_eqt,ip.n_var,ip.p,y_collo,ip.Bx_mean,ip.By0,ip.B_0,V_A,ip.mu,ip.k_x,Lm,ip.tao_r,ip.tao_d,phi40_contra)
        # H1,f1 = ps3.assign(ip.n_eqt,ip.n_var,U0,y_collo,ip.Bx_mean,ip.B_0,V_A,ip.mu,ip.k_x,Lm,ip.tao_r,ip.tao_d,phi40_contra)
        H1,f1 = ps4.assign(ip.n_eqt,ip.n_var,U0,y_collo,ip.Bx_mean,V_A,ip.mu,ip.k_x,Lm,ip.tao_r,ip.tao_d,phi40_contra)
        a1 = linalg.solve(H1,f1)

        # #?--------------- doppler shift -----------------------------------
        # for k in range(3):
        #     z = a1[ip.n_order*2:ip.n_order*3-1][k*2]
        #     doppler_shift[i,k,0] = np.abs(z)
        #     doppler_shift[i,k,1] = np.angle(z)
        
        #?----------------------------------------------------------------
        # kelvin_i[i] = a1[0] # kelvin mode intensity #?--------------------
        #?----------------------------------------------------------------

        #?-------------------------eigenfrequency-------------------------
        # R = ps4.eigv(ip.n_p,ip.n_eqt,ip.n_var,phi40_contra)
        # iw_raw = linalg.eigvals(a =H1,b = 1j*R)
        # iw_sort = np.sort(iw_raw.real)
        # iw = iw_sort[(iw_sort >=-ip.k_x) & (iw_sort <=(ip.k_x + 1e-4))] # extract value for physical mode
        # rossby1_10[:5,i] = iw[2:11:2]
        # rossby1_10[5,i] = iw[-1]
        #?----------------------------------------------------------------


        

        u = np.zeros((ip.Ny,ip.Nx),dtype=complex)
        v,eta,bx,by = [np.copy(u)]*4
        
        x = np.linspace(-ip.x_nondim,ip.x_nondim,ip.Nx)
        y = np.linspace(-ip.y_nondim,ip.y_nondim,ip.Ny)
        # y = [0.0]

        phi40_grid_contra = test2.phi_grid_array(y,ip.n_order)
        u,v,eta,bx,by = ps4.assign_3d(ip.n_var,ip.k_x,ip.n_order,a1,x,phi40_grid_contra)


        # #! ---------------------hotspot------------------------------#
        # equtor_index = int((ip.Ny-1)/2)
        # # equtor_index = 0
        # tmp_max[i] = x[np.argmax(np.real(eta)[equtor_index,:])]

        # #! -----------------acceleration curve-----------------------#
        # acc_term = acc.acc_term(u,v,eta,bx,by,y,x,V_A,Lm,ip.tao_r,ip.tao_d,ip.k_x)

        # #! ------------------acccurve data to .csv-----------------------#
        # data_acccurve = pd.DataFrame({'y':y,'term1':acc_term[0],'term2':acc_term[1],'term3':acc_term[2],'term4':acc_term[3],'term5':acc_term[4]})
        # data_acccurve.to_csv('../data/acc_curve_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'.csv',index=False,sep=',')

        # #!-------------------contour data to .npy------------------------#
        np.save('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_u.npy',u)
        np.save('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_v.npy',v)
        np.save('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_eta.npy',eta)
        np.save('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_bx.npy',bx)
        np.save('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_by.npy',by)


        
    # width_Lm[j,:] = tmp_max

    print('Lm = '+'%.5f'%(Lm)+'  completed, LOOP '+str(j))
    # print(rossby1_10)

    # #! ---------------------hotspot data to .csv--------------------#
    # dataframe = pd.DataFrame({'V_A':V_A_list,'maximum':tmp_max})

    # !-------------------- Kelvin wave intensity ----------------------#
    # np.save('../data/mode_split/kelvin/'+'Lm'+'%.3f'%(Lm)+'.npy',kelvin_i)

    #!---------------------eigenfrequency------------------------------#
    # np.save('../data/mode_split/eigenfreq/'+'Lm'+'%.3f'%(Lm)+'_2.npy',rossby1_10)

    #!---------------------doppler shift-------------------------------#
    # np.save('../../data/mode_split/doppler/'+'Lm'+'%.3f'%(Lm)+'.npy',doppler_shift)

    # #将DataFrame存储为csv,index表示是否显示行名，default=True
    # dataframe.to_csv('../data/hotspot_0408/hs_V_A'+'Lm'+'%.3f'%(Lm)+'_revised2.csv',index=False,sep=',')
    j = j+1


    
 
