import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import netCDF4

#Case study instant file
nc = netCDF4.Dataset ('/bdd/CFMIP/workspace/rguzman/GOCCP_v3.0/run.CFMIP2_200706night/out/instant/instant_SR_CR_DR_2007-06-15T20-42-00ZN_night_CFMIP2_3.0.nc')
 
# recupere la variable latitude, SR et OPAC_FA
lat = nc.variables['latitude'][:]
inst_cloud = nc.variables['instant_Cloud_OPAQ'][:].T
inst_opaq = nc.variables['instant_OPAQ'][:].T
alt = nc.variables['alt_mid'][:]

nc.close()

#Case study extent, profiles #
beg = 15440
end = 15600

#Define the discrete colormap Cloud Mask
nb_features1 = 7
cmaplist1 = [np.array((0.,0.,0.)) for i in np.arange(nb_features1)] 
cmaplist1[0] = np.array((40.,40.,40.))/255    # Nan
cmaplist1[1] = np.array((250.,250.,250.))/255     # Clear
cmaplist1[2] = np.array((40.,60.,230.))/255    # Cloud
cmaplist1[3] = np.array((50.,150.,200.))/255    # Uncertain
cmaplist1[4] = np.array((220.,220.,50.))/255    # Surface
cmaplist1[5] = np.array((100.,100.,100.))/255    # Rejected
cmaplist1[6] = np.array((240.,30.,30.))/255    # FA
my_cmap1 = mpl.colors.ListedColormap(cmaplist1)
colorbins1 = [0.5, 1.5, 2.5, 4.5, 5.5, 6.5, 7.5, 8.5]
my_norm1 = mpl.colors.BoundaryNorm(colorbins1, my_cmap1.N)

#Define the discrete colormap OPAQ Mask
nb_features2 = 9
cmaplist2 = [np.array((0.,0.,0.)) for i in np.arange(nb_features2)] 
cmaplist2[0] = np.array((40.,40.,40.))/255    # Default
cmaplist2[1] = np.array((40.,60.,230.))/255    # Cloud
cmaplist2[2] = np.array((250.,250.,250.))/255    # Clear
cmaplist2[3] = np.array((50.,150.,200.))/255    # Uncertain
cmaplist2[4] = np.array((220.,220.,50.))/255    # Weak_signal
cmaplist2[5] = np.array((180.,70.,30.))/255    # Clear_FA
cmaplist2[6] = np.array((180.,120.,30.))/255    # Uncert_FA
cmaplist2[7] = np.array((240.,30.,30.))/255    # FA
cmaplist2[8] = np.array((40.,220.,40.))/255    # z_opaque
my_cmap2 = mpl.colors.ListedColormap(cmaplist2)
colorbins2 = [-0.5, 0.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5, 10.5]
my_norm2 = mpl.colors.BoundaryNorm(colorbins2, my_cmap2.N)

#FIGURE
fig = plt.figure(figsize=(10,10))

#Cloud mask
plt.subplot(2,1,1)
cf1 = plt.pcolor(lat[beg:end], alt, inst_cloud[:,beg:end], cmap=my_cmap1, norm=my_norm1)
cb1 = plt.colorbar(cf1)
cb1.set_ticks([0.5, 1.5, 2.5, 4.5, 5.5, 6.5, 7.5, 8.5])
cb1.set_ticklabels([' ']) # Delete colorbar number label
for j, lab in enumerate(['Nan', 'Clear', 'Cloud', 'Uncertain', 'Surface','Rejected','FA_v.2']):
    cb1.ax.text(1.2, 1/14.0 + j/7.0, lab, va='center')

plt.ylabel('Altitude (km)')
plt.axis([lat[end], lat[beg], alt.min(), alt.max()])
plt.title('a) GOCCP v3 Cloud mask, case study: Cirrus + convection', fontsize=14)

#OPAQ mask
plt.subplot(2,1,2)
cf2 = plt.pcolor(lat[beg:end], alt, inst_opaq[:,beg:end], cmap=my_cmap2, norm=my_norm2)
cb2 = plt.colorbar(cf2)
cb2.set_ticks([-0.5, 0.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5, 10.5])
cb2.set_ticklabels([' ']) # Delete colorbar number label
for j, lab in enumerate(['Default', 'Cloud', 'Clear', 'Uncertain', 'Weak_signal', 'Clear_FA', 'Uncert_FA', 'FA', 'z_opaque']):
    cb2.ax.text(1.2, 1/18.0 + j/9.0, lab, va='center')

plt.xlabel('Latitude (degrees N)')
plt.ylabel('Altitude (km)')
plt.axis([lat[end], lat[beg], alt.min(), alt.max()])
plt.title('b) GOCCP v3 OPAQ mask, case study: Cirrus + convection', fontsize=14)

#save the figure
plt.savefig('Instant_masks_2007-06-15T20-42-00ZN.png')

#end