#!/usr/bin/env python

import xarray as xr
import matplotlib.pyplot as plt
import numpy as np

#  period 1989-2005  (index 1)
#fileT1='t0n110w_5day_1989-2005.cdf'
#  period 2003-2007  (index 2)
#fileT1='t0n110w_5day_2003-2007.cdf'
#  period 1989-2019  (index 2)
fileT1='t0n110w_5day_1989-2019.cdf'
fileU1='cur0n110w_5day_1989-2019.cdf'


dsT1=xr.open_dataset(fileT1)
TT1=dsT1.T_20.values
TT1=np.squeeze(np.squeeze(TT1,axis=2),axis=2)
TT1.shape
latT1=dsT1.lat.values
lonT1=dsT1.lon.values
depthT1=dsT1.depth.values
timeT1=dsT1.time.values
for ii in range(depthT1.shape[0]) :
   print(ii,depthT1[ii])

dsU1=xr.open_dataset(fileU1)
UU1=dsU1.U_320.values
UU1=np.squeeze(np.squeeze(UU1,axis=2),axis=2)
UU1[UU1==1e+35]=np.nan
UU1.shape
latU1=dsU1.lat.values
lonU1=dsU1.lon.values
depthU1=dsU1.depth.values
timeU1=dsU1.time.values
for ii in range(depthU1.shape[0]) :
   print(ii,depthU1[ii])

#pylab


#  TT1 1 pas de temps en plus que UU1!!!
Timedim=timeU1.shape[0]


TT2=np.ones((Timedim,5))*np.nan
UU2=np.ones((Timedim,5))*np.nan

UU2[:,0:5]=UU1[:,1:6]
TT2[:,0]=TT1[:-1,3]
TT2[:,1]=TT1[:-1,7]
TT2[:,2]=TT1[:-1,13]
TT2[:,3]=TT1[:-1,19]
TT2[:,4]=TT1[:-1,22]

TT3=np.ones((Timedim,5))*np.nan
UU3=np.ones((Timedim,5))*np.nan

compteur=0
for ii in range(Timedim) :
#for ii in range(10) :
# print (ii,timeT1[ii],timeU1[ii])
# if not(np.any(np.isnan(UU2[ii,:]))) :
  if (not(np.any(np.isnan(UU2[ii,:])))) & (not(np.any(np.isnan(TT2[ii,:])))) :
     print(ii,timeT1[ii],UU2[ii,:],TT2[ii,:])
#    print(ii,timeT1[ii],UU2[ii,:])
#    print (ii)
     UU3[ii,:]=UU2[ii,:]
     TT3[ii,:]=TT2[ii,:]
     compteur+=1
print (compteur,' valeurs completes sur ',Timedim)

legend1=[]

# 1989-2005    t0n110w_5day_1989-2005.cdf
#for kk in [0,15,16,19,21,22,23,24] :

# 2003-2007     t0n110w_5day_2003-2007.cdf'
#for kk in [0,11,12,14,15,17,18,19] :

# 1989-2019     t0n110w_5day_1989-2019.cdf'
for kk in [1,3,7,13,19,22,26] :
#for kk in range(depthT1.shape[0]) :
   TT1a=TT1[:,kk]
   TT1a_mean=TT1a[~np.isnan(TT1a)].mean()
   TT1a_min=TT1a[~np.isnan(TT1a)].min()
   TT1a_max=TT1a[~np.isnan(TT1a)].max()
   label1='THETAO (C)    %6.0fm    %8.1f C'%(depthT1[kk], TT1a_mean)
   print(label1, '     min: %8.1f    max: %8.1f'%(TT1a_min,TT1a_max))
   plt.figure(kk)
  #legend1.append(plt.plot(TT1a,label=label1))
   legend1=plt.plot(TT1a,label=label1)
   plt.plot([0,Timedim],[TT1a_mean,TT1a_mean])
   plt.text(Timedim,TT1a_mean,'%-4.1f'%(TT1a_mean),bbox=dict(facecolor='red'),alpha=0.5,fontsize=20,fontweight='bold')
   plt.plot([0,Timedim],[0,0])
   plt.grid()
   plt.xlim(0,Timedim)
   plt.ylim(0,30)
   plt.legend(loc='upper right')
   plt.title('TAO  (110W/0N)  1989-2019') 
   plt.savefig('TAO_110W_0N_thetao_1989-2019_%.0fm_k_%d.png'%(depthT1[kk],kk))
#  plt.show()
   plt.close(kk)

print ('\n\n')

#for kk in range(depthU1.shape[0]) :
for kk in [0,1,2,3,4,5,6] :
   UU1a=UU1[:,kk]
   UU1a_mean=UU1a[~np.isnan(UU1a)].mean()
   UU1a_min=UU1a[~np.isnan(UU1a)].min()
   UU1a_max=UU1a[~np.isnan(UU1a)].max()
   label1='UO (cm/s)    %6.0fm    %8.0f cm/s'%(depthU1[kk], UU1a_mean)
   print(label1, '     min: %8.2f    max: %8.2f'%(UU1a_min,UU1a_max))
   plt.figure(kk)
  #legend1.append(plt.plot(UU1a,label=label1))
   legend1=plt.plot(UU1a,label=label1)
   plt.plot([0,Timedim],[UU1a_mean,UU1a_mean])
   plt.text(Timedim,UU1a_mean,'%-4.0f'%(UU1a_mean),bbox=dict(facecolor='green'),alpha=0.5,fontsize=20,fontweight='bold')
   plt.plot([0,Timedim],[0,0])
   plt.grid()
   plt.xlim(0,Timedim)
   plt.ylim(-200,200)
   plt.legend(loc='upper right')
   plt.title('TAO  (110W/0N)  1989-2019') 
   plt.savefig('TAO_110W_0N_uo_1989-2019_%.0fm_k_%d.png'%(depthU1[kk],kk))
#  plt.show()
   plt.close(kk)



#############################    SELECTION DES VALEURS COMPLETES SUR TOUTE LA VERTICALES
# 1989-2019     t0n110w_5day_1989-2019.cdf'
for kk in range(TT3.shape[1]) :
   TT3a=TT3[:,kk]
   TT3a_mean=TT3a[~np.isnan(TT3a)].mean()
   TT3a_min=TT3a[~np.isnan(TT3a)].min()
   TT3a_max=TT3a[~np.isnan(TT3a)].max()
   label1='THETAO (C)    %6.0fm    %8.1f C'%(depthT1[kk], TT3a_mean)
   print(label1, '     min: %8.1f    max: %8.1f'%(TT3a_min,TT3a_max))
   plt.figure(kk)
  #legend1.append(plt.plot(TT3a,label=label1))
   legend1=plt.plot(TT3a,label=label1)
   plt.plot([0,Timedim],[TT3a_mean,TT3a_mean])
   plt.text(Timedim,TT3a_mean,'%-4.1f'%(TT3a_mean),bbox=dict(facecolor='red'),alpha=0.5,fontsize=20,fontweight='bold')
   plt.plot([0,Timedim],[0,0])
   plt.grid()
   plt.xlim(0,Timedim)
   plt.ylim(0,30)
   plt.legend(loc='upper right')
   plt.title('TAO  (110W/0N)  1989-2019') 
   plt.savefig('complete_TAO_110W_0N_thetao_1989-2019_k_%d.png'%(kk))
#  plt.show()
   plt.close(kk)

print ('\n\n')

for kk in range(UU3.shape[1]) :
   UU3a=UU3[:,kk]
   UU3a_mean=UU3a[~np.isnan(UU3a)].mean()
   UU3a_min=UU3a[~np.isnan(UU3a)].min()
   UU3a_max=UU3a[~np.isnan(UU3a)].max()
   label1='UO (cm/s)    %6.0fm    %8.0f cm/s'%(depthU1[kk], UU3a_mean)
   print(label1, '     min: %8.2f    max: %8.2f'%(UU3a_min,UU3a_max))
   plt.figure(kk)
  #legend1.append(plt.plot(UU3a,label=label1))
   legend1=plt.plot(UU3a,label=label1)
   plt.plot([0,Timedim],[UU3a_mean,UU3a_mean])
   plt.text(Timedim,UU3a_mean,'%-4.0f'%(UU3a_mean),bbox=dict(facecolor='green'),alpha=0.5,fontsize=20,fontweight='bold')
   plt.plot([0,Timedim],[0,0])
   plt.grid()
   plt.xlim(0,Timedim)
   plt.ylim(-200,200)
   plt.legend(loc='upper right')
   plt.title('TAO  (110W/0N)  1989-2019') 
   plt.savefig('complete_TAO_110W_0N_uo_1989-2019_k_%d.png'%(kk))
#  plt.show()
   plt.close(kk)



#dsT1.close()
#dsU1.close()