# coding: utf8
#*****************************************************
# PROGRAMME         : SM ANALYSIS
# AUTEUR            : C. MAGAND adapted from B. DELORME
# CREATION          : 05/09/2016
# PYTHON            : PYTHON2.7
# OBJECTIF          : Comparaison des humidités du sol - simu CMIP6 - OBS CCI-ESA v02.2
#******************************************************

#**** Packages loading
import netCDF4 
import scipy 
import numpy as np
import datetime as dt
import calendar
import pandas as pd
import pratic_climato as p
import matplotlib.cm as cm
import calendar as cd
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
import numpy.ma as ma
import math

# period data
periodlgth=1 #30
nmonths=12*periodlgth

# zones data
zones=['Boreal_North_America','Temperate_North_America','Europe','Boreal_Asia','Temperate_Asia','Tropical_Asia','Australia_NewZealand','North_Africa','South_Africa','Tropical_South_America','South_America_Temperate','Global']
zones=['Boreal_North_America','Temperate_North_America','Europe','Boreal_Asia','Temperate_Asia','Global']
#zones=['Boreal North America','Temperate North America','Europe','Boreal Asia','Temperate Asia','Global']
#zones=['Global','Boreal_North_America','Boreal_Asia','Temperate_North_America','Europe','Temperate_Asia']
ind_zones=[0,1,2,3,4,15]
#ind_zones=[15,0,3,1,2,4]
#zones=['Northern Europe','Central Mid-Europe','Eastern Europe','Mediterranean Europe','Boreal Asian Forest','Boreal Asian Steppe',
#'Boreal Asian tundra','South America temperate','Canada','Europe Union 15']
#ind_zones=[31,32,33,34,39,40,41,52,54,58]
#zones=['Australia','Tropical Asia','Northern Europe','Boreal Asia','Temperate Asia','Arabia','Temperate Asia East','Tundra']
#ind_zones=[6,25,31,35,36,37,38,53]
nb_zones=len(ind_zones) #12
print nb_zones

#obs_file=indir+'fco2_JENA-s81-V3-6_Oct2014-ext3_1981-2000_monthlymean_XYT_144x143_TS.nc'
#mod_file=indir+'NPv5.70PDctrl01_19530101_19721231_1M_nbp_time_TS.nc'

#mod_file1=indir+'IPSL-CM5A-MR_19860101_20051231_1M_nbp_time_TS.nc'
#mod_file2=indir+'CM605-LR-pdCtrl-01_20600101_20791231_1M_nbp_time_TS.nc'
#mod_file3=indir+'CM605.calv-LR-pdCtrl-02_20900101_21091231_1M_nbp_time_TS.nc'
#mod_file4=indir+'CM605.dt20-LR-pdCtrl-02_20500101_20691231_1M_nbp_time_TS.nc'
#mod_file5=indir+'CM605.GUST-LR-pdCtrl-01_20800101_20991231_1M_nbp_time_TS.nc'
#mod_file6=indir+'CM605.NOSU-LR-pdCtrl-03_19800101_19991231_1M_nbp_time_TS.nc'
#mod_file7=indir+'CM605.THC1-LR-pdCtrl-01_21000101_21191231_1M_nbp_time_TS.nc'
#mod_file8=indir+'CM605.Z0-LR-pdCtrl-01_20800101_20991231_1M_nbp_time_TS.nc'
#mod_file9=indir+'NPv5.70PDctrl01_19530101_19721231_1M_nbp_time_TS.nc'
#mod_file10=indir+'v5.71vd2qsp375_21020101_21211231_1M_nbp_time_TS.nc'
#mod_file11=indir+'FG1low.Choi.CN3238.RUN.NewX_19940101_20131231_1M_nbp_time_TS.nc'


#*** Input files and directories
noun='CMIP6'
#path='/data/cpipsl/CMIP6/INPUT/nbp/CMIP6-IPSL/CLIMATO'
path='/data/cpipsl/BCOM/INPUT/CMIP5/LAND/nbp/r1i1p1'
path2='/data/cpipsl/BCOM/INPUT/INVERSIONS'
#path='/data/cpipsl/CMIP6/INPUT/nbp/CMIP6-IPSL/SEMI-VARIOGRAM'  #'/Users/cmagand/pro/CMIP6/IGCM_OUT/'
#simus=['IPSL-CM5A-MR','CM605-LR-pdCtrl-01','CM605.calv-LR-pdCtrl-02','CM605.dt20-LR-pdCtrl-02','CM605.GUST-LR-pdCtrl-01','CM605.NOSU-LR-pdCtrl-03',
#'CM605.THC1-LR-pdCtrl-01','CM605.Z0-LR-pdCtrl-01'] # Attention : ne pas modifier l'ordre des simulations. + rajouter simus L Fairhead.
#'fco2_JENA-s81-V3-6_Oct2014-ext3',

#simus=['IPSL-CM5A-MR_shiftlonlatold','fco2_JENA-s81-V3-6_Oct2014-ext3','CM605.calv-LR-pdCtrl-02','CM605.dt20-LR-pdCtrl-02'],'CM605.GUST-LR-pdCtrl-01','CM605-LR-pdCtrl-01','CM605.NOSU-LR-pdCtrl-03','CM605.THC1-LR-pdCtrl-01','CM605.Z0-LR-pdCtrl-01','NPv5.70PDctrl01','v5.71vd2qsp375']
#simus=['BNU-ESM','CanESM2','CCSM4','CESM1-BGC','CESM1-CAM5','CESM1-FASTCHEM','CESM1-WACCM','CNRM-ESM1','GFDL-ESM2G','GFDL-ESM2M','HadGEM2-CC','HadGEM2-ES','inmcm4','IPSL-CM5A-LR','IPSL-CM5A-MR','IPSL-CM5B-LR','MIROC-ESM-CHEM','MIROC-ESM','MPI-ESM-LR','MPI-ESM-MR','MPI-ESM-P','MRI-ESM1','NorESM1-ME','NorESM1-M']
simus=['fco2_JENA-s81-V3-6_Oct2014-ext3', 'CESM1-BGC','IPSL-CM5A-MR','NorESM1-ME', 'OCN_trendy_S3']

#,'CM605.GUST-LR-pdCtrl-01','CM605-LR-pdCtrl-01','CM605.NOSU-LR-pdCtrl-03','CM605.THC1-LR-pdCtrl-01','CM605.Z0-LR-pdCtrl-01','NPv5.70PDctrl01','v5.71vd2qsp375']
#simus=['fco2_JENA-s81-V3-6_Oct2014-ext3','IPSL-CM5A-MR','CM605-LR-pdCtrl-01','CM605.calv-LR-pdCtrl-02',  'CM605.dt20-LR-pdCtrl-02','CM605.GUST-LR-pdCtrl-01','CM605.NOSU-LR-pdCtrl-03', 'CM605.THC1-LR-pdCtrl-01','CM605.Z0-LR-pdCtrl-01','NPv5.70PDctrl01','v5.71vd2qsp375']#'IPSL-CM5A-MR','CM605-LR-pdCtrl-01',
nsimus=len(simus)
#suffix=['19500101_21091231','19500101_20691231','19500101_20991231','19500101_20191231','19500101_21391231','19500101_20991231']
#1981-2000',
#'1981-2000',
suffix=['19860101_20051231','1981-2000', '20800101_20991231','19530101_19721231','21020101_21211231']#,'20600101_20791231','19800101_19991231','21000101_21191231','20800101_20991231','19530101_19721231','21020101_21211231']
#suffix=['1981-2000','19860101_20051231','20600101_20791231','20900101_21091231','20500101_20691231','20800101_20991231','19800101_19991231','21000101_21191231','20800101_20991231','19530101_19721231','21020101_21211231'] #,'19500101_20691231','19500101_20991231','19500101_20191231','19500101_21391231','19500101_20991231']

#mask_05 = "/Users/cmagand/pro/CMIP5/TOOLS/regions_mask_extend_land_143x144.nc"
#mask_05='/data/cpipsl/CMIP6/INPUT/nbp/CMIP6-IPSL/SEMI-VARIOGRAM'+'/regions_mask_extend_land_oct2015_144x143.nc'
mask_05='/data/cpipsl/BCOM/TOOLS' + '/regions_mask_extend_land_oct2015_96x73.nc'

#obs_dir='/Users/cmagand/pro/DATA/SM/ESACCI-SOILMOISTURE-L3S-SSMV-COMBINED_v02.2-1978-2014-fv02.2/monthly/2deg/'
#obs_file=obs_dir+'ESACCI-SOILMOISTURE-L3S-SSMV-COMBINED-%(year)s-2deg-monthly_means_5obsmin-fv02.2.nc'

# Coordinates
filelab=path+'/'+simus[2]+'_historical.mon.land.Lmon.r1i1p1.nbp_1980_2004_96x73_cdo_ymonmean.nc'
#filelab=path+'/'+simus[1]+'_'+suffix[1]+'_monthlymean_XYT_144x143_ymonmean_units.nc'


#filelab=path+simus[1]+'/'+simus[1]+'_'+suffix[1]+'_1M_nbp_time.nc'
print(filelab)
nc = netCDF4.Dataset(filelab,"r")
lons = nc.variables["lon"][:]
lats = nc.variables["lat"][:]
lon_ = len(nc.dimensions["lon"])
lat_ = len(nc.dimensions["lat"])
nc.close()

filelab='/data/cpipsl/BCOM/TOOLS/CONTFRAC_96x73.nc'
# Areas of cells
nc= netCDF4.Dataset(filelab,"r")
#areas=nc.variables["Areas"][:]
frac=nc.variables["CONTFRAC"][:]
nc.close()

filelab= "/data/cpipsl/BCOM/TOOLS/area_for_96x73.nc"
nc= netCDF4.Dataset(filelab,"r")
areas=nc.variables["AREA"][:]
nc.close()

areacell=areas*frac
areacell=areacell.filled(0)

# Graphs parameters

ymin=-10
ymax=10
cmap= cm.get_cmap('terrain_r')
units_lab='PgC/mth'
col=['dodgerblue','red','forestgreen','cyan','coral','brown','yellow','magenta']

model_cmip5 =["BNU-ESM","CanESM2","CCSM4","CESM1-BGC","CESM1-CAM5","CESM1-FASTCHEM","CESM1-WACCM","GFDL-ESM2G","GFDL-ESM2M","HadGEM2-CC","HadGEM2-ES","inmcm4","IPSL-CM5A-LR","IPSL-CM5A-MR","IPSL-CM5B-LR","MIROC-ESM","MIROC-ESM-CHEM","MPI-ESM-LR","MPI-ESM-MR","MPI-ESM-P","MRI-ESM1","NorESM1-M","NorESM1-ME","CNRM-ESM1"]
#col_cmip5 =["#C55CA2", "#AAACAE", "#3B63E6", "#865CA5", "#FF4500","#FFE000", "#34D314", "#75CFE5", "#FAB641", "#3DB555", "#6DBE50", "#FF1392","#474096", "#00B3C3", "#FFF02D", "#F04F38", "#3DB555", "#474096", "#13BABA","#FFA400", "#775BB4", "#C55CA2", "#F5E6BD","#B12121","#3B63E6", "#FF4500", "#FFE000","#34D314", "#FF1392", "#FFA400", "#775BB4", "#F5E6BD", "#789BF1", "#AFDFE6","#0000C7", "#1FB1AA","#DEB886", "#FFB6C1"]
col_cmip5 =["black", "#865CA5", "#00B3C3", "#DEB886", "red", "#AAACAE", "#3B63E6", "red", "#FFE000", "#34D314", "#75CFE5", "#FAB641", "#3DB555", "#6DBE50", "#FF1392","#474096", "#FFF02D", "#F04F38", "#3DB555", "#474096", "#13BABA","#FFA400", "#775BB4", "#C55CA2", "#F5E6BD","#B12121","#3B63E6", "#FF4500", "#FFE000","#34D314", "#FF1392", "#FFA400", "#775BB4", "#F5E6BD", "#789BF1", "#AFDFE6","#0000C7", "#1FB1AA", "#FFB6C1"]
#col_cmip5 =["black", "#AAACAE", "#3B63E6", "#865CA5", "red","#FFE000", "#34D314", "#75CFE5", "#FAB641", "#3DB555", "#6DBE50", "#FF1392","#474096", "#00B3C3", "#FFF02D", "#F04F38", "#3DB555", "#474096", "#13BABA","#FFA400", "#775BB4", "#C55CA2", "#F5E6BD","#B12121","#3B63E6", "#FF4500", "#FFE000","#34D314", "#FF1392", "#FFA400", "#775BB4", "#F5E6BD", "#789BF1", "#AFDFE6","#0000C7", "#1FB1AA","#DEB886", "#FFB6C1"]
col=[col_cmip5[13],col_cmip5[0:nsimus]]
lsty_cmip5=['solid','solid','solid','solid','solid','solid','solid','solid','solid','solid','solid','solid','solid',
'dashed','dashed','dashed','dashed','dashed','dashed',
'solid','solid','dashed','solid','solid']
lsty=[lsty_cmip5[13],lsty_cmip5[0:nsimus]]
#*** Lecture des simulations
maps=np.zeros((lat_,lon_,nsimus))
season_mod=np.zeros((nsimus,12,nb_zones))
n=0
for isimu in simus:
    print(isimu)
    #if n==0:
    #    mod_file='/Users/cmagand/pro/CMIP6/IGCM_OUT/IPSL-CM5A-MR/mrsos_Lmon_IPSL-CM5A-MR_historical_r1i1p1_185001-200512.nc'
    #    nc = netCDF4.Dataset(mod_file,"r")
    #    tmp=nc.variables["mrsos"][:,::-1,:]*0.01 # conversion en m3/m3
     ##   nc.close()
     #   lons=lons-180
    if (isimu=='fco2_JENA-s81-V3-6_Oct2014-ext3'):
       mod_file=path2+'/'+"fco2_JENA-s81-V3-6_Oct2014-ext3_1981-2004_monthlymean_XYT_96x73_cdo_ymonmean.nc"
       #mod_file=path+'/'+simus[n]+'_'+suffix[n]+'_monthlymean_XYT_144x143_ymonmean_units.nc'
       nc = netCDF4.Dataset(mod_file,"r")
#       tmp=nc.variables["nbp"][:]*0.01 # conversion en m3/m3
       tmp=nc.variables['Terrestrial_flux'][:] /12./1E15 #/12
       nc.close()
    elif (isimu=='OCN_trendy_S3'):
#       mod_file1='/data/cpipsl/CMIP6/INPUT/nbp/CMIP6-IPSL/CLIMATO/CM605.calv-LR-pdCtrl-02_20900101_21091231_1M_nbp_time_ymonmean_units.nc'
#       nc1 = netCDF4.Dataset(mod_file1,"r")
#       tmp1=nc1.variables["nbp_units"][:]/1E15 #/12
#       nc1.close()
       mod_file='/data/cpipsl/CMIP6/INPUT/nbp/CLIMATO/Ncycle/OCN_trendy_S3_historical.land.nbp_1979_2009_96x73_ymonmean.nc'
       nc = netCDF4.Dataset(mod_file,"r")
#       tmp=nc.variables["nbp"][:]*0.01 # conversion en m3/m3
       tmp=nc.variables['nbp'][:]*1000.*24.*365./12./1.E15*(-1)
#       tmp[tmp==0] = np.NaN
#       print tmp[9][59]
       #tmp=tmp/1E15 
       #print tmp[9][59]
#       ma.masked_where(tmp == -0.000000, tmp)
       nc.close()
#       tmp_mask = tmp.mask
#       print tmp[9][59]

       tmp=tmp.filled(np.NaN)
    else:
       mod_file=path+'/'+simus[n]+'_historical.mon.land.Lmon.r1i1p1.nbp_1980_2004_96x73_cdo_ymonmean.nc'
   #    mod_file=path+simus[n]+'/'+simus[n]+'_'+suffix[n]+'_1M_nbp_time.nc'

       nc = netCDF4.Dataset(mod_file,"r")
  #    tmp=nc.variables["nbp"][:]*0.01 # conversion en m3/m3
       tmp=nc.variables['nbp'][:] /1.E15 *86400*365/12*1000*(-1)
 #      tmp=nc.variables['nbp'][:] *86400*365/12*(-1)
       #print tmp[9][59]
       nc.close()
    
    print(mod_file)
    ntimesimu=tmp.shape[0]
    mod_data=tmp[(ntimesimu-nmonths):ntimesimu,:,:]
    #print mod_data.shape
    #****** Calcul du cycle saisonnier pour cette simulation et sur les 12 régions TRANSCOM
    print("seasonal cycle")
    season_mod[n,:,:]=p.seasonal_cycle(mod_data,areacell,mask_05,nb_zones,zones,ind_zones,noun,nmonths,isimu)
    #mod_tmp=np.zeros((12,mod_data.shape[1],mod_data.shape[2]))
    #for t in range(12):
    #    mod_tmp[t,:,:]=np.nanmean(mod_data[range(t,nmonths,12),:,:],axis=0)
        #mod=np.nanmean(mod_season,axis=(1,2))
    #season_mod[:,n]=np.nanmean(mod_tmp,axis=(1,2))#/np.nansum(areacell)
    #season_mod[:,n]=np.nansum(mod_tmp*areacell,axis=(1,2))/np.nansum(areacell)          
    #****** Plot de la carte d'humidité moyenne
    print("maps")
    maps[:,:,n]=np.nanmean(mod_data,axis=0)*1E15
    #print maps[45,45,:]
    #print maps.shape
    #print areacell[45,45]
    #print areacell.shape

    lon, lat = np.meshgrid(lons,lats)
    p.mean_map(maps[:,:,n],lat,lon,noun,isimu,areacell,ymin,ymax,cmap,units_lab)
    #p.season_maps()
    #****** Plot des histogrammes en fonction des régions
    print("histo")
    #p.histo()
    n+=1

# Seasonal cycles for all simulations and all regions
fig, ax = plt.subplots(3,2,figsize=(20,17))
#fig, ax = plt.subplots(1,1)

fig.suptitle(noun+' seasonal_cycle ', fontsize=18)
for i, ax in enumerate(ax.flat,start=1):
    zo=i-1
    print(zones[zo])
    xtick=np.arange(0,12.5,1)
    x=[0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5, 10.5, 11.5] #np.arange([0.5,11.5,1])
    monthnum = (np.linspace(1,12,num=12))
    mons=cd.month_abbr[1:13]
    #print monthnum
#    text_file = open("/data/sli/PC/seasonalcycle_22sept_climato/output/seasonal_cycle_"+zones[zo]+".txt", "w")

    for j in range(nsimus):
        #if (j==0):
        #    ax.plot(x,season_mod[j,:,zo],color=col[j],linestyle=lsty[j],label=simus[j],lw=2)
        #else:
        ax.plot(x,season_mod[j,:,zo],color=col_cmip5[j],linestyle=lsty_cmip5[j],label=simus[j],lw=2.5)
        #text_file.write(simus[j]+zones[zo]+": %s" % season_mod[j,:,zo])
#    text_file.write("%5s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s" % ("month", simus[0][0:6], simus[1][0:6], simus[2][0:6], simus[3][0:6], simus[4][0:6], simus[5][0:6], simus[6][0:6], simus[7][0:6], simus[8][0:6], simus[9][0:6], simus[10][0:6], simus[11][0:6], simus[12][0:6], simus[13][0:6], simus[14][0:6], simus[15][0:6], simus[16][0:6], simus[17][0:6], simus[18][0:6], simus[19][0:6], simus[20][0:6], simus[21][0:6], simus[22][0:6], simus[23][0:6]))
#    text_file.write('\r\n')
    #for imon in range(0,12):   
#        text_file.write("%5s  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f  %10.7f"  % (monthnum[imon], season_mod[0,imon,zo], season_mod[1,imon,zo], season_mod[2,imon,zo], season_mod[3,imon,zo], season_mod[4,imon,zo], season_mod[5,imon,zo], season_mod[6,imon,zo], season_mod[7,imon,zo], season_mod[8,imon,zo], season_mod[9,imon,zo], season_mod[10,imon,zo], season_mod[11,imon,zo], season_mod[12,imon,zo], season_mod[13,imon,zo], season_mod[14,imon,zo], season_mod[15,imon,zo], season_mod[16,imon,zo], season_mod[17,imon,zo], season_mod[18,imon,zo], season_mod[19,imon,zo], season_mod[20,imon,zo], season_mod[21,imon,zo], season_mod[22,imon,zo], season_mod[23,imon,zo]))
 #       text_file.write("%5s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s  %s"  % (monthnum[imon], season_mod[0,imon,zo], season_mod[1,imon,zo], season_mod[2,imon,zo], season_mod[3,imon,zo], season_mod[4,imon,zo], season_mod[5,imon,zo], season_mod[6,imon,zo], season_mod[7,imon,zo], season_mod[8,imon,zo], season_mod[9,imon,zo], season_mod[10,imon,zo], season_mod[11,imon,zo], season_mod[12,imon,zo], season_mod[13,imon,zo], season_mod[14,imon,zo], season_mod[15,imon,zo], season_mod[16,imon,zo], season_mod[17,imon,zo], season_mod[18,imon,zo], season_mod[19,imon,zo], season_mod[20,imon,zo], season_mod[21,imon,zo], season_mod[22,imon,zo], season_mod[23,imon,zo]))
#        text_file.write('\r\n')
#    text_file.write('\r\n')
    #np.savetxt('/home/wflmd/PC/seasonalcycle_PC/Jorge/output/seasonal_cycle.txt', season_mod[j,:,zo], delimiter=',')   # X is an array 
    ax.set_xticks(xtick)
    ax.set_xticklabels(mons)
    ax.set_title(zones[zo],fontsize=25)
    ax.grid()
    ax.set_xlabel('month')
    ax.set_ylabel('PgC/mth')
#    text_file.close()
#    plt.xlabel("month")
#    plt.ylabel(units_lab)
    #if (i==9):
    #if (i==6):
    ax.legend(loc='best',prop={'size':6})
plt.subplots_adjust(left=0.15)

#plt.savefig('/Users/cmagand/pro/CMIP6/GRAPHS/CMIP6/seasonal_cycle_'+noun+'_allsimulations_regions.png')
#plt.savefig('/home/wflmd/PC/seasoncycle/output/seasonal_cycle_'+noun+'_allsimulations_regions.png')
#plt.savefig('/data/sli/PC/cyclesaisonnier_climato/output/seasonal_cycle_climato_'+noun+'_allsimulations_'+zones[zo]+'.png')
#plt.savefig('/data/cpipsl/CMIP6/ANALYSIS/cyclesaisonnier_climato/output/seasonal_cycle_climato_'+noun+'_allsimulations_'+zones[zo]+'.png')
plt.savefig('/data/sli/PC/seasonalcycle_22sept_climato/output/seasonal_cycle_climato_'+noun+'_allsimulations.png')
#plt.savefig('/data/sli/PC/seasonalcycle_22sept_climato/output/seasonal_cycle_climato_'+noun+'_a    llsimulations.png')
plt.close()


# Histogram
#bins = np.linspace(0,0.5,50)
#print "histogram"
#p.histo(map_obs,map_mod,mask_05,nb_zones,zones,ind_zones,bins,noun,yvec)
