import argparse
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import netCDF4
from scipy import fftpack

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


#---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.29, 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
 


u = np.load('../0npy/u_av.npy',allow_pickle=True)
w = np.load('../0npy/w_av.npy',allow_pickle=True)
theta = np.load('../0npy/theta_av.npy',allow_pickle=True)




z_Votx = av(av(u[:,:,:],1),1)  #t,z,y,x
x_Votx = av(av(u[:,:,:],0),0) 
y_Votx = av(av(u[:,:,:],0),1)
 
zz = np.linspace(0,4000,len(z_Votx))/20000
xx = np.linspace(-50000,50000,len(x_Votx))/20000


X,Z = np.meshgrid(xx,zz)

theta_avt = av(theta[:,50000/200:50600/200,:],1)
vari = theta_avt
vari_min = np.min(theta_avt)
vari_max = np.max(theta_avt)
print("vari_min, max",vari_min ,vari_max)
vari_min = 290
vari_max = 320
u =av(u[:,50000/200:50600/200,:],1)#/U_ref #* 100
w =av(w[:,50000/200:50600/200,:],1)#/U_ref* 100





#imj = plt.streamplot(X,Z,u,w, color= 'k',linewidth=2, density=1.5,cmap='jet',arrowsize = 2.5)





vari[vari==0.0] = nan
masked_analysis= np.ma.array(vari, mask=np.isnan(vari))
cmap = matplotlib.cm.jet
cmap.set_bad('gray', alpha=1.0)
im = plt.pcolormesh(X,  Z, masked_analysis,vmin=vari_min, vmax=vari_max, cmap=cmap, shading='gouraud', edgecolors='None', alpha=1.0)


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


w_sub[:,20:] = 0.0

# Create the quiver plot with reduced vectors


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





plt.xlabel(r'$x/D$')
plt.ylabel(r'$z/D$')
yticks = np.arange(0, 0.2, 0.05)
xticks = np.arange(-2.5, 3.5,1)
ax.set_xticks(xticks)
ax.set_yticks(yticks)
plt.ylim([0, 0.1])
plt.xlim([-2.5, 2.5])

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,330,10))
cbar_label =r'$ \theta \ \mathrm{(K)}$'   
cbar.ax.set_ylabel(cbar_label, labelpad=1.0)



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



u = np.load('../0npy/u_av.npy',allow_pickle=True)
v = np.load('../0npy/v_av.npy',allow_pickle=True)



h1=h1/40
h2=h2/40


u = u[h1:h2,:,:]
v = v[h1:h2,:,:]

print("u",np.shape(u))

z_Votx = av(av(u[:,:,:],1),1)  #t,z,y,x
x_Votx = av(av(u[:,:,:],0),0) 
y_Votx = av(av(u[:,:,:],0),1)
 
zz = np.linspace(0,4000,len(z_Votx))/4000
xx = np.linspace(-50000,50000,len(x_Votx))/20000


X,Z = np.meshgrid(xx,xx)

u_avt =av(u,0)
v_avt = av(v,0)
vari = np.sqrt(u_avt**2 + v_avt**2)
vari =  vari/4.73

u =av(u,0)#/U_ref #* 100
v =av(v,0)#/U_ref* 100




print("!!!",np.max(u[:]))



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

arrowstyle='-|>'


vari[vari==0.0] = nan
masked_analysis= np.ma.array(vari, mask=np.isnan(vari))
cmap = matplotlib.cm.jet
cmap.set_bad('gray', alpha=1.0)
im = plt.pcolormesh(X,  Z, masked_analysis,vmin=0, vmax=0.8, cmap=cmap, shading='gouraud', edgecolors='None', alpha=1.0)
ax.set_aspect(1)

xx1 = 100  
xx2 = 400  


#imj = plt.streamplot(X[xx1:xx2,xx1:xx2],  Z[xx1:xx2,xx1:xx2],u[xx1:xx2,xx1:xx2],v[xx1:xx2,xx1:xx2], color= 'k',linewidth=1.5, density=density, arrowsize =arrowsize ,cmap='jet')


####################
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)



step =12
plt.quiver(X[100:701:step,100:701:step],Z[100:701:step,100:701:step],u[100:701:step,100:701:step],v[100:701:step,100:701:step], headwidth=9, headlength=15, pivot='mid' )


plt.xlabel(r'$x/D$',fontsize=26)
plt.ylabel(r'$y/D$',fontsize=26)
yticks = np.arange(-2.5, 3.5, 0.5)
xticks = np.arange(-2.5, 3.5,0.5)
ax.set_xticks(xticks)
ax.set_yticks(yticks)
plt.ylim([-1.0, 1.0])
plt.xlim([-1.0, 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,35,0.2))
#cbar_label =r'$ \sqrt{ \overline{u}_x^{2} + \overline{u}_y^{2} }  $ / $   U_D $'   
cbar_label =r'$ \sqrt{u_x^{2} + u_y^{2} }  $ / $   U_D $'
cbar.ax.set_ylabel(cbar_label,fontsize=26,labelpad=3.0)

#ax.get_xaxis().set_visible(False)
#ax.get_yaxis().set_visible(False)


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







h1=h1/40
h2=h2/40

print("h1",h1)
u =uu[h1:h2,:,:]
v =vv[h1:h2,:,:]
w =ww[h1:h2,:,:]

theta =thetat[h1:h2,:,:]



z_Votx = av(av(u[:,:,:],1),1)  #t,z,y,x
x_Votx = av(av(u[:,:,:],0),0) 
y_Votx = av(av(u[:,:,:],0),1)



 

xx = np.linspace(-50000,50000,len(x_Votx))/20000


X,Z = np.meshgrid(xx,xx)

u_avt = av(u,0)
v_avt = av(v,0)
vari = np.sqrt(u_avt**2 + v_avt**2)


theta_av =av(theta,0)



theta_av =  (theta_av - thetaa)/(Delta_thetam) 


vari_min =  np.min(theta_av)
vari_max = np.max(theta_av)
print("vari_min, max",vari_min ,vari_max)

vari[vari==0.0] = nan
masked_analysis= np.ma.array(theta_av, mask=np.isnan(vari))
cmap = matplotlib.cm.jet
cmap.set_bad('gray', alpha=1.0)
im = plt.pcolormesh(X,  Z, 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('./30xy'+name+'.png')