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

#parameters
version = '3.1.2'
dayflag = 'night'

#Tables
CF3D_cloud     = ma.ones((40,90,180,31))*-9999. # (alt,lat,lon,day)
CF3D_thick     = ma.ones((40,90,180,31))*-9999.
CF3D_thin      = ma.ones((40,90,180,31))*-9999.
CF3D_opacity   = ma.ones((40,90,180,31))*-9999.

#Fill value
CF3D_cloud.set_fill_value([-9999.])
CF3D_thick.set_fill_value([-9999.])
CF3D_thin.set_fill_value([-9999.])
CF3D_opacity.set_fill_value([-9999.])

year = 0
annee = str(year + 2010)
mois = '06'

#June 2010
liste = glob.glob('/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP_v3/3D_CloudFraction/grid_2x2xL40/'+annee+'/'+dayflag+'/daily/3D_CloudFraction_OPAQ330m_'+annee+mois+'??_'+dayflag+'_CFMIP2_sat_3.1.2.nc')

#Reading all the files from the list
for i_fic in np.arange(len(liste)):
#Opening GOCCP daily file
    nc = netCDF4.Dataset (liste[i_fic])
#Reading the variables
    CF3D_thick[:,:,:,i_fic] = nc.variables['clcalipso_opaque'][0][:]
    CF3D_thin[:,:,:,i_fic]  = nc.variables['clcalipso_notopaque'][0][:]
    CF3D_opacity[:,:,:,i_fic] = nc.variables['calipsoopacity'][0][:]

    if i_fic == 0 :
        lon = nc.variables['longitude'][:]
        lat = nc.variables['latitude'][:]
        alt = nc.variables['alt_mid'][:]
        alti_b = nc.variables['alt_bound'][:]
        alt_b = alti_b[0]

    nc.close()

#Masking fill values
CF3D_thick = ma.masked_where(CF3D_thick < 0, CF3D_thick)
CF3D_thin = ma.masked_where(CF3D_thin < 0, CF3D_thin)
CF3D_opacity = ma.masked_where(CF3D_opacity < 0, CF3D_opacity)
#Computing the All cloud variable
CF3D_cloud = CF3D_thick + CF3D_thin
CF3D_cloud = ma.masked_where(CF3D_cloud < 0, CF3D_cloud)

#Monthly means
CF3D_cloud_mean = ma.mean(CF3D_cloud, axis=3)
CF3D_thick_mean = ma.mean(CF3D_thick, axis=3)
CF3D_thin_mean = ma.mean(CF3D_thin, axis=3)
CF3D_opacity_mean = ma.mean(CF3D_opacity, axis=3)
#Zonal mean
CF3D_cloud_zonal = ma.mean(CF3D_cloud_mean, axis=2)
CF3D_thick_zonal = ma.mean(CF3D_thick_mean, axis=2)
CF3D_thin_zonal = ma.mean(CF3D_thin_mean, axis=2)
CF3D_opacity_zonal = ma.mean(CF3D_opacity_mean, axis=2)

#MASK fractions < 1%
cloud_mask = ma.masked_where(CF3D_cloud_zonal < 0.01, CF3D_cloud_zonal)
thick_mask = ma.masked_where(CF3D_thick_zonal < 0.01, CF3D_thick_zonal)
thin_mask = ma.masked_where(CF3D_thin_zonal < 0.01, CF3D_thin_zonal)
opacity_mask = ma.masked_where(CF3D_opacity_zonal < 0.01, CF3D_opacity_zonal)

#FIGURE
fig = plt.figure(figsize=(10,8))
fig.subplots_adjust(bottom=0.08, left=0.1, top = 0.9, right=0.95)
fig.suptitle('GOCCP_v'+version+' zonal Cloud and Opacity fractions, '+mois+' '+annee+' '+dayflag, fontsize=16)

plt.subplot(221)
cf1 = plt.pcolormesh(lat, alt_b, cloud_mask*100., vmin=1, vmax=0.3*100, cmap='CMRmap_r')
cb1 = plt.colorbar(cf1)
plt.ylabel('Altitude (km)')
plt.xlim(-82,82)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-80, 81, 20), ['80S', '60S', '40S', '20S', '0', '20N', '40N', '60N', '80N'], fontsize=10)
plt.title("a) All cloud (%)",fontsize=12)

plt.subplot(222)
cf2 = plt.pcolormesh(lat, alt_b, thick_mask*100, vmin=1, vmax=0.3*100, cmap='CMRmap_r')
cb2 = plt.colorbar(cf2)
plt.ylabel('Altitude (km)')
plt.xlim(-82,82)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-80, 81, 20), ['80S', '60S', '40S', '20S', '0', '20N', '40N', '60N', '80N'], fontsize=10)
plt.title("b) Opaque cloud profiles (%)",fontsize=12)

plt.subplot(223)
cf3 = plt.pcolormesh(lat, alt_b, thin_mask*100, vmin=1, vmax=0.3*100, cmap='CMRmap_r')
cb3 = plt.colorbar(cf3)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-82,82)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-80, 81, 20), ['80S', '60S', '40S', '20S', '0', '20N', '40N', '60N', '80N'], fontsize=10)
plt.title("c) Thin cloud profiles (%)",fontsize=12)

plt.subplot(224)
cf4 = plt.pcolormesh(lat, alt_b, opacity_mask*100, vmin=1, vmax=0.4*100, cmap='CMRmap_r')
cb4 = plt.colorbar(cf4)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-82,82)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-80, 81, 20), ['80S', '60S', '40S', '20S', '0', '20N', '40N', '60N', '80N'], fontsize=10)
plt.title("d) Opacity fraction (%)",fontsize=12)

#save the figure
plt.savefig('GOCCP_v'+version+'_CF3D_zonal_'+mois+annee+'_'+dayflag+'.png')

#end