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,8))*-9999. # (lat,lon,month,year)
map2d_zopac = ma.ones((90,180,12,8))*-9999.
map2d_zopac_mask = ma.ones((90,180))*-9999.

#Fill value
map2d_thick.set_fill_value(-9999.)
map2d_zopac.set_fill_value(-9999.)
map2d_zopac_mask.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+'/Map_OPAQ330m_'+annee+mois+'_'+day+'_CFMIP2_sat_3.1.1.nc')
#Opaque cloud cover
        map2d_thick[:,:,month,year] = nc.variables['cltcalipso_opaque'][0][:]
#z_opaque (altitude of opacity of CALIPSO)
        map2d_zopac[:,:,month,year] = nc.variables['zopaque'][0][:]

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

#Closing NCDF files
        nc.close()

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

#Annual means
map2d_thick_year = map2d_thick.mean(axis=2)
map2d_zopac_year = map2d_zopac.mean(axis=2)
#Multi-year mean
thick_clim = map2d_thick_year.mean(axis=2)
zopac_clim = map2d_zopac_year.mean(axis=2)

#Two different Masks for z_opaque
zopac000_mask = ma.masked_where(thick_clim < 0., zopac_clim)
zopac030_mask = ma.masked_where(thick_clim < 0.3, zopac_clim)

#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
#Z_opaque altitude
#All z_opaque
a_zopac000 = ma.masked_where(zopac000_mask < 0., a*thick_clim) #area and cover-weighted
zopac000_pond = zopac000_mask*a_zopac000
cczopac000 = zopac000_pond.sum()/a_zopac000.sum()
cczopac000 = float("{0:.2f}".format(cczopac000))
cczopac000_str = str(cczopac000)
#z_opaque only where opaque cloud cover > 30%
a_zopac030 = ma.masked_where(zopac030_mask < 0., a*thick_clim) #area and cover-weighted
zopac030_pond = zopac030_mask*a_zopac030
cczopac030 = zopac030_pond.sum()/a_zopac030.sum()
cczopac030 = float("{0:.2f}".format(cczopac030))
cczopac030_str = str(cczopac030)


#FIGURE
#my_cmap2 = cm.jet_r
my_cmap2 = cm.cubehelix #Thibault
#bounds2 =  [0.4, 1, 1.5, 2, 2.5, 3, 5, 7, 10, 13]
bounds2 =  np.arange(19)*0.5    #Altitude
norm2 = mpl.colors.BoundaryNorm(bounds2, my_cmap2.N)
fig = plt.figure(figsize=(12,4))
fig.subplots_adjust(bottom=0.10, left=0.10, top = 0.90, right=0.90)

plt.subplot(121)
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, zopac000_mask, bounds2, cmap=my_cmap2, norm=norm2)
plt.colorbar(label='Altitude (km)', orientation='horizontal')
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 z_opaque ('+cczopac000_str+' km)\nopaque cloud cover > 0, '+period, fontsize=14)

plt.subplot(122)
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, zopac030_mask, bounds2, cmap=my_cmap2, norm=norm2)
plt.colorbar(label='Altitude (km)', orientation='horizontal')
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 z_opaque ('+cczopac030_str+' km)\nopaque cloud cover > 0.3, '+period, fontsize=14)

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

#end