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
 




canyon0_x = np.arange(0,100000-200,200)/200
canyon0_y = np.arange(0,100000-200,200)/200
canyon0_z = np.arange(0,4000,50)/50




print("@@",np.shape(canyon0_x),np.shape(canyon0_y),np.shape(canyon0_z),type(canyon0_x))

 

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

T_2K = av(av(av(theta[41,:,:],0),0),0)
T_2K_H0 = av(av(av(theta[2,:,:],0),0),0)
print("Temperature at 2KM",T_2K,T_2K_H0,np.shape(theta))


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

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

 
zz = np.linspace(0,4000,len(z_Votx))/20000
xx = np.linspace(-50000,50000,len(x_Votx))/20000
X,Z = np.meshgrid(xx,zz)

plot_y = np.arange(50000,51000,200)/200
plot_y  = plot_y.round().astype(int)





print('!!!!',np.shape(theta))
theta_avt = av(theta[:,plot_y,:],1)   
P = av(P[:,plot_y,:],1)

vari = theta_avt 

#vari = (vari + 298.0 )*(1000/ (P+1000))**0.286 
vari =  vari  + 298.0 

theta_avt =  vari 



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[:,plot_y,:-1],1)#/U_ref #* 100
w =av(w[:-1,plot_y,:],1)#/U_ref* 100
print('!!!!',np.shape(u),np.shape(w))



 
#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


plt.quiver(X_sub, Z_sub, u_sub, w_sub, scale=20, headwidth=8, headlength=10, 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('./Fig4c.png')