import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from mpl_toolkits.basemap import Basemap
import matplotlib as mpl
import netCDF4

#parameters
nb_years = 8
period = '2008-2015'
#version = '3.1.1'
day = 'night' #'night' ou 'day'
#Variable arrays
map2d_thick = ma.ones((90,180,12,nb_years))*-9999. # (lat,lon,month,year)
map2d_thin = ma.ones((90,180,12,nb_years))*-9999.
map2d_clear = ma.ones((90,180,12,nb_years))*-9999.

#Fill value
map2d_thick.set_fill_value(-9999.)
map2d_thin.set_fill_value(-9999.)
map2d_clear.set_fill_value(-9999.)

for year in np.arange(nb_years):
    annee = str(year + 2008)
    for month in np.arange(12):
        mois = '{:02d}'.format(month+1)

#Reading the NCDF monthly files
#OPAQ files
        nc = netCDF4.Dataset ('/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP_v3/2D_Maps/grid_2x2xL40/'+annee+'/'+day+'/MapLowMidHigh330m_'+annee+mois+'_'+day+'_CFMIP2_sat_3.1.1.nc')
#v2.X files
        nc_opaq = netCDF4.Dataset ('/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP_v3/2D_Maps/grid_2x2xL40/'+annee+'/'+day+'/Map_OPAQ330m_'+annee+mois+'_'+day+'_CFMIP2_sat_3.1.1.nc')

#Opaque cloud cover
        map2d_thick[:,:,month,year] = nc_opaq.variables['cltcalipso_opaque'][0][:]
#Thin cloud cover
        map2d_thin[:,:,month,year] = nc_opaq.variables['cltcalipso_thin'][0][:]
#Clear sky
        map2d_clear[:,:,month,year] = nc.variables['clccalipso'][0][:]

#Reading axes
        if year ==0 and month == 0 :
            lon = nc.variables['longitude'][:]
            lat = nc.variables['latitude'][:]

        nc.close()
        nc_opaq.close()

#Masking fill values
map2d_thick = ma.masked_where(map2d_thick < 0., map2d_thick)
map2d_thin = ma.masked_where(map2d_thin < 0., map2d_thin)
map2d_clear = ma.masked_where(map2d_clear < 0., map2d_clear)

#Annual means
map2d_thick_year = map2d_thick.mean(axis=2)
map2d_thin_year = map2d_thin.mean(axis=2)
map2d_clear_year = map2d_clear.mean(axis=2)
#Multi-year mean
thick_clim = map2d_thick_year.mean(axis=2)
thin_clim = map2d_thin_year.mean(axis=2)
clear_clim = map2d_clear_year.mean(axis=2)

#Compute global mean values
#meshgrid with latitudes and longitudes 
m = Basemap(projection='cyl', llcrnrlon=-180, urcrnrlon=180, llcrnrlat=-90, urcrnrlat=90, resolution='c')
lons, lats = np.meshgrid(lon,lat)
x, y = m(lons,lats)
a = np.cos((y*2*np.pi)/360) #tableau des surfaces ponderees
#Opaque cloud cover
a_thick = ma.masked_where(thick_clim < 0., a) #area-weighted
thick_pond = thick_clim*a_thick
ccopaque = int(round(100.*thick_pond.sum()/a_thick.sum())) # fraction --> % (integer)
ccopaque_str = str(ccopaque)
#Thin cloud cover
a_thin = ma.masked_where(thin_clim < 0., a)
thin_pond = thin_clim*a_thin
ccthin = int(round(100.*thin_pond.sum()/a_thin.sum()))
ccthin_str = str(ccthin)
#Clear sky
a_clear = ma.masked_where(clear_clim < 0., a)
clear_pond = clear_clim*a_clear
ccclear = int(round(100.*clear_pond.sum()/a_clear.sum()))
ccclear_str = str(ccclear)


#FIGURE
#my_cmap = cm.RdYlBu_r
my_cmap = cm.Reds
bounds =  np.arange(11)*0.1
norm = mpl.colors.BoundaryNorm(bounds, my_cmap.N)
fig = plt.figure(figsize=(8,10))
fig.subplots_adjust(bottom=0.05, left=0.1, top = 0.95, right=0.97)

plt.subplot(311)
m = Basemap(projection='cyl', llcrnrlon=-180, urcrnrlon=180, llcrnrlat=-90, urcrnrlat=90, resolution='c')
m.drawcoastlines(color='k')
lons, lats = np.meshgrid(lon,lat)
x, y = m(lons,lats)
m.contourf(x, y, thick_clim, bounds, cmap=my_cmap, norm=norm)
plt.colorbar(label='Cover (Fraction)')
plt.xticks(range(-180, 181, 60), ['180W', '120W', '60W', '0E', '60E', '120E', '180E'], fontsize=10)
plt.ylim(-82, 82)
plt.yticks([-80, -60, -30, 0, 30, 60, 80], ['80S', '60S', '30S', '0', '30N', '60N', '80N'], fontsize=10)
plt.title('a) GOCCP v3.0 opaque cloud cover ('+ccopaque_str+'%), '+period, fontsize=14)

plt.subplot(312)
m = Basemap(projection='cyl', llcrnrlon=-180, urcrnrlon=180, llcrnrlat=-90, urcrnrlat=90, resolution='c')
m.drawcoastlines(color='k')
lons, lats = np.meshgrid(lon,lat)
x, y = m(lons,lats)
m.contourf(x, y, thin_clim, bounds, cmap=my_cmap, norm=norm)
plt.colorbar(label='Cover (Fraction)')
plt.xticks(range(-180, 181, 60), ['180W', '120W', '60W', '0E', '60E', '120E', '180E'], fontsize=10)
plt.ylim(-82, 82)
plt.yticks([-80, -60, -30, 0, 30, 60, 80], ['80S', '60S', '30S', '0', '30N', '60N', '80N'], fontsize=10)
plt.title('b) GOCCP v3.0 thin cloud cover ('+ccthin_str+'%), '+period, fontsize=14)

plt.subplot(313)
m = Basemap(projection='cyl', llcrnrlon=-180, urcrnrlon=180, llcrnrlat=-90, urcrnrlat=90, resolution='c')
m.drawcoastlines(color='k')
lons, lats = np.meshgrid(lon,lat)
x, y = m(lons,lats)
m.contourf(x, y, clear_clim, bounds, cmap=my_cmap, norm=norm)
plt.colorbar(label='Cover (Fraction)')
plt.xticks(range(-180, 181, 60), ['180W', '120W', '60W', '0E', '60E', '120E', '180E'], fontsize=10)
plt.ylim(-82, 82)
plt.yticks([-80, -60, -30, 0, 30, 60, 80], ['80S', '60S', '30S', '0', '30N', '60N', '80N'], fontsize=10)
plt.title('c) GOCCP v3.0 clear sky cover ('+ccclear_str+'%), '+period, fontsize=14)

#save the figure
plt.savefig('GOCCP_v3.0_covers_'+period+'.png')

#end