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

CF3D_cloud     = ma.ones((40,90,180,12,8))*-9999. # (alt,lat,lon,month,year)
CF3D_thick     = ma.ones((40,90,180,12,8))*-9999.
CF3D_thin      = ma.ones((40,90,180,12,8))*-9999.
CF3D_zopac     = ma.ones((40,90,180,12,8))*-9999.
CF3D_zopac_sum = ma.ones((40,90,180,12,8))*-9999.

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

for year in np.arange(8):
    annee = str(year + 2008)

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

#Reading the NCDF monthly files
        nc = netCDF4.Dataset ('/bdd/CFMIP/workspace/rguzman/GOCCP_v3.0/run.CFMIP2_'+annee+mois+'night/out/mensuel/3D_CloudFraction330m_'+annee+mois+'_night_CFMIP2_sat_3.0.nc')

#Reading the different variables
        CF3D_cloud[:,:,:,month,year] = nc.variables['clcalipso'][0][:]
        CF3D_thick[:,:,:,month,year] = nc.variables['clcalipso_opaque'][0][:]
        CF3D_thin[:,:,:,month,year]  = nc.variables['clcalipso_notopaque'][0][:]
        CF3D_zopac[:,:,:,month,year] = nc.variables['calipsozopaque'][0][:]

#Reading the axes once
        if year == 0 and month == 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()

#Computing the opacity profiles
#Vertical sum of the z_opaque fractions
#From the highest level (39) ...
CF3D_zopac_sum[39,:,:,:,:] = CF3D_zopac[39,:,:,:,:]
#... to the surface (38 --> 0)
for z in np.arange(39):
    CF3D_zopac_sum[38-z,:,:,:,:] = CF3D_zopac_sum[39-z,:,:,:,:] + CF3D_zopac[38-z,:,:,:,:]

#Annual means
CF3D_cloud_year = ma.mean(CF3D_cloud, axis=3)
CF3D_thick_year = ma.mean(CF3D_thick, axis=3)
CF3D_thin_year = ma.mean(CF3D_thin, axis=3)
CF3D_zopac_year = ma.mean(CF3D_zopac, axis=3)
CF3D_zopac_sum_year = ma.mean(CF3D_zopac_sum, axis=3)
#Multi-year means
CF3D_cloud_mean = ma.mean(CF3D_cloud_year, axis=3)
CF3D_thick_mean = ma.mean(CF3D_thick_year, axis=3)
CF3D_thin_mean = ma.mean(CF3D_thin_year, axis=3)
CF3D_zopac_mean = ma.mean(CF3D_zopac_year, axis=3)
CF3D_zopac_sum_mean = ma.mean(CF3D_zopac_sum_year, axis=3)
#Zonal means
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_zopac_zonal = ma.mean(CF3D_zopac_mean, axis=2)
CF3D_zopac_sum_zonal = ma.mean(CF3D_zopac_sum_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)
zopac_mask = ma.masked_where(CF3D_zopac_zonal < 0.01, CF3D_zopac_zonal)
zopac_sum_mask = ma.masked_where(CF3D_zopac_sum_zonal <= 0.01, CF3D_zopac_sum_zonal)

#FIGURE
fig = plt.figure(figsize=(10,8))
fig.subplots_adjust(bottom=0.06, left=0.1, top = 0.95, right=0.95)

plt.subplot(221)
cf1 = plt.pcolormesh(lat, alt_b, cloud_mask*100., vmin=1, vmax=0.2*100)
cb1 = plt.colorbar(cf1)
plt.ylabel('Altitude (km)')
plt.xlim(-90,90)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-90, 91, 30), ['90S', '60S', '30S', '0', '30N', '60N', '90N'], fontsize=10)
plt.title("a) All cloud (%), 2008-2015",fontsize=12)

plt.subplot(222)
cf2 = plt.pcolormesh(lat, alt_b, thick_mask*100, vmin=1, vmax=0.1*100)
cb2 = plt.colorbar(cf2)
plt.ylabel('Altitude (km)')
plt.xlim(-90,90)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-90, 91, 30), ['90S', '60S', '30S', '0', '30N', '60N', '90N'], fontsize=10)
plt.title("b) Opaque cloud profiles (%), 2008-2015",fontsize=12)

plt.subplot(223)
cf3 = plt.pcolormesh(lat, alt_b, thin_mask*100, vmin=1, vmax=0.1*100)
cb3 = plt.colorbar(cf3)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-90,90)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-90, 91, 30), ['90S', '60S', '30S', '0', '30N', '60N', '90N'], fontsize=10)
plt.title("c) Thin cloud profiles (%), 2008-2015",fontsize=12)

plt.subplot(224)
cf4 = plt.pcolormesh(lat, alt_b, zopac_sum_mask*100, vmin=1, vmax=0.4*100)
cb4 = plt.colorbar(cf4)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-90,90)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-90, 91, 30), ['90S', '60S', '30S', '0', '30N', '60N', '90N'], fontsize=10)
plt.title("d) Opacity fraction (%), 2008-2015",fontsize=12)

#save the figure
plt.savefig('GOCCP_v3.0_CF3D_zonal_opacity_2008-2015.png')

#end