'''The following script plots annual mean of cloud area fraction for the 4 different satellite products from ICDC for a given year. The scripts takes circa. 3 minutes to plot data and requires Python >= version 3.10 with numpy, matplotlib, xarray, netcdf4, cartopy and dask. You may also need to update your dask diagnostics, e.g. via "python -m pip install dask[diagnostics] --upgrade".'''

import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
import cartopy.crs as ccrs

### --- src functions --- ###
def figure_setup(title):
  ''' returns figure with 5 axes that use the Robinson projection
  and have coastlines.'''

  fig = plt.figure(figsize=(12, 6))
  axs = []
  a=1
  for i in range(3):
    for j in range(2):
      axs.append(fig.add_subplot(2, 3, a, projection=ccrs.Robinson()))
      a+=1
  fig.delaxes(axs[-1])
  for ax in axs:
    ax.coastlines()

  fig.suptitle(title)

  return fig, axs[:-1]

def plot_annualmean_on_axis(ax, axis_title, cbar_label, dataset_path, var_name,
                            scale_factor, cmap='plasma', vmin=0, vmax=100):
  '''creates a single dataset from all the (netdcf) files found in the 'dataset_path',
  then takes the time mean of the variable called 'var_name' in the dataset, and
  multiplies it by a scale_factor. Finally plots the mean values as an image
  on the given axis 'ax' with title 'axis_title'.'''
  ds = xr.open_mfdataset(dataset_path)

  ntimes = ds.time.size
  print('Mean over '+str(ntimes)+' times in datasets from '+dataset_path)
  yearmean = ds[var_name].mean('time') * scale_factor

  cbar_kwargs={'shrink': 0.5, 'label': cbar_label}
  yearmean.plot(ax=ax, transform=ccrs.PlateCarree(), cmap=cmap,
                cbar_kwargs=cbar_kwargs, vmin=vmin, vmax=vmax)
  ax.set_title(axis_title)

def save_figure(fig, savefig_name='figure1.png', is_show=True):
  ''' saves 'fig' as file called 'savefig_name'.'''

  fig.savefig(savefig_name, dpi=400, bbox_inches="tight", facecolor='w', format="png")
  print("Figure saved as: "+savefig_name)
  
  if is_show:
    plt.show()
### --------------------- ###

### --- input parameters --- ###
year = '2010' # year to plot (must be available in all 5 products)
fig_title = 'Mean of Total Cloud Fraction in '+year
cbar_label = 'Annual Mean Cloud Fraction /%'
savefig_name = "./cloud_area_fraction.png"

icdcpath = '/pool/data/ICDC/atmosphere/'
basepaths = { # base of paths to netcdf files to use in plotting of each satellite product
  'calipso_cloudsat': icdcpath+'calipso_cloudsat_cloudcover/DATA/',
  'clara3': icdcpath+'eumetsat_clara3_cloud/DATA/',
  'isccp': icdcpath+'isccp/DATA/',
  'modis_aqua': icdcpath+'modis_aqua_cloud/DATA/',
  'modis_terra': icdcpath+'modis_terra_cloud/DATA/',
}

datasets = { # full paths to netcdf files to use in plotting of each satellite product
    'calipso_cloudsat': basepaths['calipso_cloudsat']+'UCAR__merged_cloudsat-calipso_totalcloudcover__UCAM-ICDC__v0.1__2deg__'+year+'*.nc',
    'clara3': basepaths['clara3']+'FractionalCloudCover/global/MONTHLY/CFCmm'+year+f'*01000000319AVPOS01GL.nc',
    'isccp': basepaths['isccp']+'hgg/'+year+'*/ISCCP-Basic.HGG.v01r00.GLOBAL.'+year+'.*.01.*.GPC.10KM.CS00.EA1.00.nc', 
    'modis_aqua': basepaths['modis_aqua']+year+'/MODIS-C6.1__MYD08__daily__cloud-Fractions__'+year+'*__UHAM-ICDC__fv0.1.nc',
    'modis_terra': basepaths['modis_terra']+year+'/MODIS-C6.1__MOD08__daily__cloud-Fractions__'+year+'*__UHAM-ICDC__fv0.1.nc',
  }

dataset_names = { # names of each satellite product
  'calipso_cloudsat': 'CALIPSO + CloudSat',
  'clara3': 'CLARA-A3',
  'isccp': 'ISCCP',
  'modis_aqua': 'MODIS AQUA',
  'modis_terra': 'MODIS TERRA',
}

variable_names = { # names of cloud area fraction in each satellite product
  'calipso_cloudsat': 'cf',
  'clara3': 'cfc',
  'isccp': 'cldamt',
  'modis_aqua': 'cfrac',
  'modis_terra': 'cfrac',
}

variable_scalefactor = { # multiply variable by scale_factor for each satellite to convert to %
  'calipso_cloudsat': 100,
  'clara3': 1,
  'isccp': 1,
  'modis_aqua': 100,
  'modis_terra': 100,
}
### ------------------------ ###

### --- plotting of data --- ###
datasets2plt = ['calipso_cloudsat', 'clara3', 'isccp', 'modis_aqua', 'modis_terra']
fig, axs = figure_setup(fig_title)
for ax, d in zip(axs, dataset_names.keys()):
  plot_annualmean_on_axis(ax, dataset_names[d], cbar_label, datasets[d],
                          variable_names[d], variable_scalefactor[d])  
fig.tight_layout()
save_figure(savefig_name)
### ------------------------ ###