import argparse
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import scipy as sp
import scipy.stats as stats
import numpy as np
from numpy import *
import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.patches as patches
import scipy.interpolate as interpolate


import argparse
parser = argparse.ArgumentParser(description='calculate mean tracer age')
parser.add_argument('-location','--location',help='project name',required=True)



args = parser.parse_args()
location = args.location



#---using Latex
font = {'family' : 'serif', 'size' : 52}
matplotlib.rc('font', **font)
plt.rc('legend', fontsize=52)
plt.rc('legend', frameon=False)
plt.rc('text',usetex=True)

fig, ax = plt.subplots(figsize=(12.0,5))
plt.subplots_adjust(left=0.18, bottom=0.28, top=0.91,right =0.84)


from numpy import nanmean
def av (arr, axis):

        arr = np.isfinite(arr)*arr
        temp_arr = nanmean(arr,axis=axis)
        temp_arr = np.isfinite(temp_arr)*temp_arr
        return temp_arr




X=np.load('../0cluster/2npy/X'+location+'.npy')
Y_ori=np.load('../0cluster/2npy/Z'+location+'.npy')


U_x=np.load('../0cluster/2npy/U_x'+location+'.npy')
U_y=np.load('../0cluster/2npy/U_z'+location+'.npy')




T_origi=np.load('../0cluster/2npy/T'+location+'.npy')
P_origi=np.load('../0cluster/2npy/P'+location+'.npy')
print("P_origi",P_origi[:10])
Theta = T_origi *(1013.25/(P_origi + 1013.25))**(287.05/1006)









Y =Y_ori

 

print(np.shape(U_x),np.shape(Y),Y.max())


density = [2,2] #[5,5] 
arrowsize = 1
arrowstyle = '->'
scale = 0.12
arrow_angles = 'xy'
scale_units = 'xy'


xi = np.linspace(X.min(), X.max(), 1000)
#print("xi",xi)
yi = np.linspace(Y.min(), Y.max(), 100)
#print(yi)
XX, YY = np.meshgrid(xi, yi)
U = interpolate.griddata((X, Y), U_x, (XX, YY), method='linear')
V = interpolate.griddata((X, Y), U_y, (XX, YY), method='linear')
 
Theta = interpolate.griddata((X, Y), Theta, (XX, YY), method='linear')


U_plot = np.zeros((100,1000))
for h in range(100):
    U_plot[h,:] =Theta[h,:] + 0.4 *h 

print("!T_origi",np.shape(U_plot),np.min(U_plot),np.max(U_plot ))


#print("!!",U)

XX = (XX-50000)/20000
YY = (YY)/20000
#plt.streamplot(XX,YY,U,V, color='k',linewidth=2, density=1.5, arrowsize =2.5 )
#print(XX)







##############

vari = U_plot

vari[vari==0.0] = nan
masked_analysis= np.ma.array(U_plot, mask=np.isnan(vari))
cmap = matplotlib.cm.jet
cmap.set_bad('gray', alpha=1.0)
im = plt.pcolormesh(XX,YY, masked_analysis,vmin=290 ,vmax=320, cmap=cmap, shading='gouraud', edgecolors='None', alpha=1.0)


step = 3  # Adjust this value to control vector density (e.g., 5, 10, 20)
stepz = 30
X_sub = XX[::step, ::stepz]
Z_sub = YY[::step, ::stepz]
u_sub = U[::step, ::stepz]  # Replace u with your actual data
w_sub = V[::step, ::stepz]  # Replace w with your actual data

# Create the quiver plot with reduced vectors


plt.quiver(X_sub, Z_sub, u_sub, w_sub, scale=20, headwidth=8, headlength=10, pivot='mid')






ax.set_yticks(np.arange(0,0.21,step=0.05))
ax.set_xticks(np.arange(-2.5,2.6,step=1.))
plt.ylim([0, 0.1])
plt.xlim([-2.5, 2.5])
plt.xlabel(r'$x/D$',fontdict=font)
plt.ylabel(r'$z/D$',fontdict=font,labelpad=1.0)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="3%", pad=0.05)
cbar=plt.colorbar(im,cax=cax, ticks= np.arange(290,340,10))
cbar_label =r'$ \theta \ \mathrm{(K)}$'   
cbar.ax.set_ylabel(cbar_label,labelpad=1.0)




plt.savefig('./Fig4b.png')




Theta = T_origi *(1013.25/(P_origi + 1013.25))**(287.05/1006)




U_x =  (T_origi - thetaa)/(Delta_thetam) 




#U_x =  (Theta - theta0)/(thetam- theta0)
print("U_x",U_x)
#print("U_x_origi",U_x_origi)

#Z = 1/0.0001 * np.log(1+0.0001 *Z_ori)





density = [2,2] #[5,5] 
arrowsize = 2
arrowstyle = '->'
scale = 0.12
arrow_angles = 'xy'
scale_units = 'xy'


xi = np.linspace(X.min(), X.max(), 1000)
#print("xi",xi)
yi = np.linspace(Y.min(), Y.max(), 1000)
XX, YY = np.meshgrid(xi, yi)
U = interpolate.griddata((X, Y), U_x, (XX, YY), method='linear')


#print("!!",U)

XX = (XX-50000)/20000
YY = (YY-50000)/20000


#print(XX)







##############

vari = U

vari[vari==0.0] = nan
masked_analysis= np.ma.array(U, mask=np.isnan(vari))
cmap = matplotlib.cm.jet
cmap.set_bad('gray', alpha=1.0)
im = plt.pcolormesh(XX,YY, masked_analysis,vmin=-0.2, vmax=0.9, cmap=cmap, shading='gouraud', edgecolors='None', alpha=1.0)
ax.set_aspect(1)
####################
xline = np.linspace(-0.5,0.5)
yline = np.linspace(0.5,0.5)
plt.plot(xline,yline,color='white',linestyle='--', linewidth=2)

xline = np.linspace(-0.5,0.5)
yline = np.linspace(-0.5,-0.5)
plt.plot(xline,yline,color='white',linestyle='--', linewidth=2)

xline = np.linspace(-0.5,-0.5)
yline = np.linspace(-0.5,0.5)
plt.plot(xline,yline,color='white',linestyle='--', linewidth=2)

xline = np.linspace(0.5,0.5)
yline = np.linspace(-0.5,0.5)
plt.plot(xline,yline,color='white',linestyle='--', linewidth=2)


ax.set_yticks(np.arange(-1.5,1.6,step=0.5))
ax.set_xticks(np.arange(-1.5,1.6,step=0.5))
plt.ylim([-1.0, 1.0])
plt.xlim([-1.0, 1.0])
plt.xlabel(r'$x/D$',fontdict=font)
plt.ylabel(r'$y/D$',fontdict=font,labelpad=1.0)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="3%", pad=0.05)
cbar=plt.colorbar(im,cax=cax, ticks= np.arange(-0.2,1,0.2))
cbar_label =r'$ ( \theta -  \theta_a ) / ( \theta_m -  \theta_0)$'   
cbar.ax.set_ylabel(cbar_label,fontsize=26,labelpad=-1.0)


plt.savefig('./xy'+location+'.png')