import numpy as np
import matplotlib.pyplot as plt
import scipy.special as spe
from scipy import linalg
import pseudo_spectral_v2 as ps2
import test2
import pandas as pd
import initial_input as ip
#!------------------------ hotspot plot ----------------------#
# plt.figure(figsize=(8,6))

# dataname = ['hs_V_ALm0.000_revised1.csv','hs_V_ALm0.079_revised1.csv','hs_V_ALm0.118_revised1.csv','hs_V_ALm0.157_revised1.csv','hs_V_ALm0.500_revised1.csv',
# 'hs_V_ALm0.768_revised1.csv','hs_V_ALm1.035_revised1.csv','hs_V_ALm1.303_revised1.csv','hs_V_ALm1.571_revised1.csv','hs_V_ALm1.771_revised1.csv',
# 'hs_V_ALm2.272_revised1.csv','hs_V_ALm2.774_revised1.csv','hs_V_ALm3.275_revised1.csv','hs_V_ALm3.776_revised1.csv','hs_V_ALm4.278_revised1.csv',
# 'hs_V_ALm4.779_revised1.csv','hs_V_ALm5.280_revised1.csv','hs_V_ALm5.782_revised1.csv','hs_V_ALm6.283_revised1.csv']

# num_of_plots = len(dataname)

# for i in range(num_of_plots):
# 	Lm = dataname[i].replace('hs_V_ALm','')
# 	Lm = Lm.replace('_revised1.csv','')
# 	width_Lm = np.array(pd.read_csv('data/'+dataname[i],skiprows = 1))
# 	#print(width_Lm)
# 	V_A_list = width_Lm[:,0]
# 	#print(V_A_list)
# 	hotspot_pos = width_Lm[:,1]
# 	c1=plt.get_cmap('RdYlBu')(np.linspace(0, 1, num_of_plots))[i]
# 	plt.plot(V_A_list,hotspot_pos,color = c1,marker = '.',label = '$L_{m}$='+Lm)
    
# plt.plot(np.linspace(0,0.2,3),[0]*3,'--',color ='black')
# plt.xticks(fontsize = 12)
# plt.yticks(fontsize = 12)
# plt.xlabel('Alfven speed $V_A$',fontsize = 15)
# plt.ylabel('Hot Spot position(x)',fontsize = 15)
# plt.xlim(0,0.25)
# #plt.ylim(-3,3)
# plt.legend()
# plt.savefig('../img/'+'width_Lm_revised2.pdf',dpi = 300)
# # plt.savefig('F:/exoplanet/Exoplanet_SWE_weno_tvd/thesis-master/image/main/sec4.2/'+'hs_V_A.pdf',dpi = 300)

# plt.show()

# #!-------------------------- contour plot ------------------------# 
x = ip.x
y = ip.y

var = ['u','v','eta','bx','by']

T = 1
Lm = np.pi
V_A_list = [0.20]

for i in range(len(V_A_list)):
	data_contour = []
	V_A = V_A_list[i]
	for j in range(len(var)):
		data_contour.append('../../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_'+var[j]+'.npy')
	u = np.real(np.load(data_contour[0])+ ip.mean_flow) 
	v = np.real(np.load(data_contour[1]))
	eta = np.real(np.load(data_contour[2])+ ip.mean_flow)
	bx = np.real(np.load(data_contour[3]))
	by = np.real(np.load(data_contour[4]))

	equtor_index = int((ip.Ny-1)/2)

	

	u_norm = u/np.max(u)
	v_norm = v/np.max(u)
	eta_norm = eta/np.max(eta)
	bx_norm = bx/np.max(by)
	by_norm = by/np.max(by)

	fig = plt.figure(figsize = (8,5))
	plt.contourf(x,y,np.real(eta_norm),20,alpha=0.9,cmap=plt.cm.hot)
	#alpha：透明度
	#cmap：color map
	n_skip_x = 10
	n_skip_y = 10
	plt.quiver(x[::n_skip_x],y[::n_skip_y],np.real(u_norm[::n_skip_x,::n_skip_y]),np.real(v_norm[::n_skip_x,::n_skip_y]),scale = 1.8,scale_units = 'inches',color = 'black',angles = 'uv',headwidth = 3,width = 0.0015)  

	#use plt.contour to add contour lines
	C=plt.contour(x,y,np.real(eta_norm),20,colors='black',alpha = 0.4)
	# plt.ylim(-np.pi,np.pi)
	#adding label
	#clabel:cycle的label，inline=True表示label在line内，fontsize表示label的字体大小
	plt.title('$V_{A}$ = '+'%.3f'%(V_A)+ ', $L_{m} = $'+ '%.3f' %(Lm)+'\n'+'$\\tau_{rad} =$'+str(ip.tao_r)+', $\\tau_{drag} =$'+str(ip.tao_d),fontsize = 15)
	plt.xlabel('Eastward distance',fontsize = 18)
	plt.ylabel('Northward distance',fontsize = 18)
	# plt.savefig('../img/revised/contourf/png/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f' %(V_A)+'_uv.png',dpi = 100)
	plt.savefig('../../img/paper/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f' %(V_A)+'_uv.pdf',dpi = 300)
	plt.show()

# #!----------------------------------acc plot--------------------------#
# V_A_list = np.linspace(0.08,0.09,11)
# Lm = np.pi

# for i in range(len(V_A_list)):
# 	V_A = V_A_list[i]
# 	acc_term = np.array(pd.read_csv('../data/acc_curve_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'.csv',skiprows = 1))
# 	y = acc_term[:,0]

# 	fig = plt.figure(figsize = (8,8))
# 	plt.plot(acc_term[:,1],y,linewidth = 2.5,color ='cyan',label = 'term1')
# 	plt.plot(acc_term[:,2],y,linewidth = 2.5,color ='black',label = 'term2')
# 	plt.plot(acc_term[:,3],y,linewidth = 2.5,color = 'darkblue',label = 'term3')
# 	plt.plot(acc_term[:,4],y,linewidth = 2.5,color = 'lightgreen',label = 'term4')
# 	plt.plot(acc_term[:,5],y,linewidth = 2.5,color = 'purple',label = 'MF')
# 	plt.plot(np.sum(acc_term[:,1:],axis = 1),y,linewidth = 2.5,color = 'red',label = 'total')
# 	# plt.xlim(-0.2,0.6)

# 	plt.plot([0.0]*5,np.linspace(-4,4,5),'--',color = 'black')
# 	plt.ylim(-ip.y_nondim,ip.y_nondim)
# 	plt.xlabel('Accelerations',fontsize = 20)
# 	plt.ylabel('y',fontsize = 20)
# 	plt.legend()
# 	plt.savefig('../img/revised/acc/png/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f' %(V_A)+'_wmf.png',dpi = 100)
# 	# plt.show()

# #! ----------------------------Hotspot - acc_max plot--------------------------#
# x = ip.x
# y = ip.y
# var = ['eta']

# T = 21
# Lm = np.pi/4
# equtor_index = int((ip.Ny-1)/2)
# V_A_list = [0.0]+list(np.linspace(0.01,0.03,21))+list(np.linspace(0.04,0.4,37))

# tmp_max = []
# acc_equator = []

# for i in range(len(V_A_list)):
# 	data_contour = []
# 	V_A = V_A_list[i]
# 	for j in range(len(var)):
# 		data_contour.append('../data/contour_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'_'+var[j]+'.npy')
	
# 	eta = np.real(np.load(data_contour[0])+ ip.mean_flow)
	
# 	tmp_max.append(x[np.argmax(eta[equtor_index,:])])

# 	acc_term = np.array(pd.read_csv('../data/acc_curve_revised/'+'Lm'+'%.3f'%(Lm)+'V_A'+'%.3f'%(V_A)+'.csv',skiprows = 1))
# 	acc_equator.append(np.sum(acc_term[equtor_index,1:]))


# fig = plt.figure(figsize = (8,8))
# plt.plot(V_A_list,tmp_max,'.-',label = 'hotspot')
# plt.plot(V_A_list,acc_equator,'.-',label = '$acc_{equ}$')
# plt.plot([0.0,0.3],[np.pi*2]*2,'--',color = 'black')
# plt.plot([0.0,0.3],[-np.pi*2]*2,'--',color = 'black')
# plt.ylabel('hotspot position & $acc_{equ}$',fontsize = 20)
# plt.xlabel('Alfven speed $V_{A}$',fontsize = 20)
# plt.xlim(0,0.3)
# plt.title('$L_{m} = \\pi/4$',fontsize  =20)
# plt.legend(fontsize = 15)
# # plt.show()
# plt.savefig('../img/revised/contra/png/'+'acc_hotspot_contra_Lm_0.785.png',dpi = 100)


#! --------------------------kelvin mode plot -------------------------------------#

# num_of_VA = 40
# # V_A_list = np.linspace(0,0.2,num_of_VA)
# 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)))
# fig = plt.figure(figsize = (8,8))
# kelvin_i = np.load('../data/mode_split/kelvin/Lm3.142.npy')

# plt.plot(V_A_list,kelvin_i,'o-')
# plt.ylabel('kelvin mode strength',fontsize = 20)
# plt.xlabel('Alfven speed $V_{A}$',fontsize = 20)
# # plt.show()
# plt.savefig('../img/revised/mode_split/png/'+'Lm_3.142.png',dpi = 100)



