import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
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.
CF3D_cloudopac = ma.ones((40,90,180,12,8))*-9999.
CF3D_cloudnotopac = ma.ones((40,90,180,12,8))*-9999.
CF3D_zopac = ma.ones((40,90,180,12,8))*-9999. # (alt,lat,lon,mois,annee)
clouds_zopac = ma.ones((40,3))*-9999.

#Fill value
CF3D_cloud.set_fill_value([-9999.])
CF3D_cloudopac.set_fill_value([-9999.])
CF3D_cloudnotopac.set_fill_value([-9999.])
CF3D_zopac.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)

#Lecture fichier NCDF
        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')

# recupere la variable latitude, SR et OPAC_FA
        CF3D_cloud[:,:,:,month,year] = nc.variables['clcalipso'][0][:]
        CF3D_cloudopac[:,:,:,month,year] = nc.variables['clcalipso_opaque'][0][:]
        CF3D_cloudnotopac[:,:,:,month,year] = nc.variables['clcalipso_notopaque'][0][:]
        CF3D_zopac[:,:,:,month,year] = nc.variables['calipsozopaque'][0][:]

# affiche les informations netCDF sur la variable, comme les unites, etc
#print latvar.units, latvar.dimensions, latvar.shape
        if (year == 0 and month == 0) :
            lon = nc.variables['longitude'][:]
            lat = nc.variables['latitude'][:]
            alt = nc.variables['alt_mid'][:]

# recupere directement les valeurs numeriques des autres variables
#lon = nc.variables['longitude'][:]
            alti_b = nc.variables['alt_bound'][:]
            alt_b = alti_b[0]
#alti_m = nc.variables['alt_mid'][:]
#alt = alti_m[:]
#T = nc.variables['ta'][0,0,:,:]

#Moyennes annuelles
CF3D_cloud_year = ma.mean(CF3D_cloud, axis=3)
CF3D_cloudopac_year = ma.mean(CF3D_cloudopac, axis=3)
CF3D_cloudnotopac_year = ma.mean(CF3D_cloudnotopac, axis=3)
CF3D_zopac_year = ma.mean(CF3D_zopac, axis=3)

#Climato, moyennes sur les annees
CF3D_cloud_mean = ma.mean(CF3D_cloud_year, axis=3)
CF3D_cloudopac_mean = ma.mean(CF3D_cloudopac_year, axis=3)
CF3D_cloudnotopac_mean = ma.mean(CF3D_cloudnotopac_year, axis=3)
CF3D_zopac_mean = ma.mean(CF3D_zopac_year, axis=3)

#Moyennes zonales
CF3D_cloud_zonal = ma.mean(CF3D_cloud_mean, axis=2)
CF3D_cloudopac_zonal = ma.mean(CF3D_cloudopac_mean, axis=2)
CF3D_cloudnotopac_zonal = ma.mean(CF3D_cloudnotopac_mean, axis=2)
CF3D_zopac_zonal = ma.mean(CF3D_zopac_mean, axis=2)

#Moyennes zone Tropical Indo-Pacific (60E-200E ; 10S-10N)
#Longitude
#cloud_prof = ma.mean(CF3D_cloud_mean, axis=2)
cloudopac_lat = ma.mean(CF3D_cloudopac_mean[:,:,120:180], axis=2)
cloudnotopac_lat = ma.mean(CF3D_cloudnotopac_mean[:,:,120:180], axis=2)
zopac_lat = ma.mean(CF3D_zopac_mean[:,:,120:180], axis=2)
#Latitude
#cloud_prof = ma.mean(CF3D_cloud_mean, axis=1)
cloudopac_prof = ma.mean(cloudopac_lat[:,40:50], axis=1)
cloudnotopac_prof = ma.mean(cloudnotopac_lat[:,40:50], axis=1)
zopac_prof = ma.mean(zopac_lat[:,40:50], axis=1)

clouds_zopac[:,0] = cloudnotopac_prof
clouds_zopac[:,1] = cloudopac_prof
clouds_zopac[:,2] = zopac_prof

#MASK fractions < 1%
cloud_mask = ma.masked_where(CF3D_cloud_zonal <= 0.001, CF3D_cloud_zonal)
cloudopac_mask = ma.masked_where(CF3D_cloudopac_zonal <= 0.001, CF3D_cloudopac_zonal)
cloudnotopac_mask = ma.masked_where(CF3D_cloudnotopac_zonal <= 0.001, CF3D_cloudnotopac_zonal)
zopac_mask = ma.masked_where(CF3D_zopac_zonal <= 0.001, CF3D_zopac_zonal)

# FIGURE 1
fig = plt.figure(figsize=(12,12))
fig.subplots_adjust(bottom=0.05, left=0.1, top = 0.95, right=0.90)
levels = [0, 1, 2, 3, 4, 6, 8, 12, 14]
cmap = plt.get_cmap('BrBG')
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)


plt.subplot(221)
#cf1 = plt.pcolormesh(lat, alt_b, np.log10(cloud_mask*100.), vmin=-1, vmax=2) #log
cf1 = plt.pcolormesh(lat, alt_b, cloud_mask*100., vmin=1, vmax=0.2*100)
cb1 = plt.colorbar(cf1)
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=12)
plt.title("a) All cloud fraction (%)",fontsize=16)

plt.subplot(222)
#cf2 = plt.pcolormesh(lat, alt_b, np.log10(ice_mask*100), vmin=-1, vmax=2) #log
cf2 = plt.pcolormesh(lat, alt_b, cloudopac_mask*100, cmap=cmap, norm=norm)
cb2 = plt.colorbar(cf2)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-60,60)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-60, 61, 20), ['60S', '40S', '20S', '0', '20N', '40N', '60N'], fontsize=12)
plt.title("a) Opaque cloud fraction (%)",fontsize=16)

plt.subplot(223)
#cf3 = plt.pcolormesh(lat, alt_b, np.log10(liq_mask*100), vmin=-1, vmax=2) #log
cf3 = plt.pcolormesh(lat, alt_b, cloudnotopac_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=12)
plt.title("c) Thin cloud fraction (%)",fontsize=16)

plt.subplot(224)
#cf4 = plt.pcolormesh(lat, alt_b, np.log10(zopac_mask*100), vmin=-1, vmax=2)#0.1*100) #log
cf4 = plt.pcolormesh(lat, alt_b, zopac_mask*100, cmap=cmap, norm=norm)
cb4 = plt.colorbar(cf4)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(-60,60)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks(range(-60, 61, 20), ['60S', '40S', '20S', '0', '20N', '40N', '60N'], fontsize=12)
plt.title("b) z_opaque fraction (%)",fontsize=16)

# sauve dans une image log
#plt.savefig('FIGURES/2008-2015/fig_3D_CloudFraction330m_GOCCP_v3.0_3zopac_zonal_sn_log10.png') ######

# sauve dans une image
#plt.savefig('FIGURES/2008-2015/fig_3D_CloudFraction330m_GOCCP_v3.0_3zopac_zonal_sn.png') ######
plt.savefig('FIGURES/2008-2015_night/CloudFraction330m_GOCCP_v3.0_zonal.png') ######

# FIGURE 2
fig = plt.figure(figsize=(5,5))
levels_2 = [0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 10, 15]
cmap_2 = plt.get_cmap('OrRd')
norm_2 = BoundaryNorm(levels_2, ncolors=cmap_2.N, clip=True)

plt.subplot(111)
#cf1 = plt.pcolormesh(lat, alt_b, np.log10(cloud_mask*100.), vmin=-1, vmax=2) #log
cf1 = plt.pcolormesh(np.arange(4), alt_b, clouds_zopac*100., cmap=cmap_2, norm=norm_2)
cb1 = plt.colorbar(cf1)
plt.xlabel('Latitude')
plt.ylabel('Altitude (km)')
plt.xlim(0,3)
plt.ylim(alt_b[0],alt_b[-1])
plt.xticks([0.5, 1.5, 2.5], ['Thin/Broken\ncloud', 'Opaque\ncloud', 'z_opaque'], fontsize=12)
plt.title("New cloud partition fractions (%)",fontsize=14)

plt.savefig('FIGURES/2008-2015_night/clouds_zopac_GOCCP_v3.0.png') ######

#fin du programme