import xarray as xr
import cftime
import matplotlib.pyplot as plt
import math
import numpy as np
ydeb=2031
yfin=2035
imon=(yfin-ydeb+1)*12
print(imon)
file_in="../../LMDZFAS4/MO/utzon_mth_"+str(ydeb)+"-"+str(yfin)+".nc"
fig_out="../../LMDZFAS4/MO/QBO_mth_"+str(ydeb)+"-"+str(yfin)+".jpg"
ds = xr.open_dataset(file_in)
vitu_4d = ds['vitu']
vitu_3d = vitu_4d.isel(lon=0)
vitu_2d = vitu_3d.sel(lat=slice(5.,-5.))
#print("vitu_2d:")
print(vitu_2d)
#print("time_var")
#print(vitu_2d.time_counter)
vitu_2d_mean=vitu_2d.mean(dim='lat')
#print("vitu_2d_mean:")
#print(vitu_2d_mean)
#vitu_2d = vitu_3d.sel(lat=0.0,method='nearest')
vitu_2d_transpose=vitu_2d_mean.transpose()
qbo=vitu_2d_transpose.values
new_time=np.linspace(1,imon,imon)
zalt=vitu_2d['presnivs'].values
new_zalt=7*(np.log(101300./zalt))
print("new_zalt:")
print(new_zalt)
niveaux = np.arange(-72.5, 72.5, 1)
nivvitu = np.arange(-52.5,52.5,5)
fig, ax = plt.subplots()
ax.set_xlim(1,imon)
ax.set_ylim(2,100000)
ax.set_xticks([0, 12, 24, 36, 48, 60,72,84,96,108,120])

#plt.gca().invert_yaxis()
plt.yscale("log")
plt.gca().invert_yaxis()
plt.contourf(new_time,zalt,qbo,levels=niveaux,cmap='seismic')
plt.colorbar()
for t in ax.get_xticks():
    ax.axvline(t, color='black', linewidth=0.9, alpha=0.5)
plt.contour(new_time,zalt,qbo,levels=nivvitu,colors='k',linewidths=0.5)
plt.title(file_in)
plt.xlabel("Mois")
plt.ylabel("z(km)")
plt.savefig(fig_out)
plt.show()
#plt.show()