#!/usr/bin/env python3
#
# Before opening Python execute:
#   module load python3/miniconda3
#   module load python3/python-rpn
#   module load python3/outils-divers
#   source activate base_plus

# Importation of modules and library
import numpy as np                         # Good module for matrix and matrix operation
import matplotlib.pyplot as plt            # Module to produce figure
import matplotlib.colors as colors
import os                                  # Used to convert png to other format

import cartopy.crs as ccrs                 # Import cartopy ccrs
import cartopy.feature as cfeature         # Import cartopy common features

import netCDF4
from datetime import datetime, timedelta
from netCDF4 import Dataset,num2date


# Example parameters which need to get adjusted !!!
# ================================================
filename = '/sca/compilers_and_tools/python/demos-python3/data/dmCAN_Rockies2_3km_P3_2013_ERAInt_201305_moyenne.nc4'    # Name of NC file to read
varname  = 'tas'    # Name of variable to read
title    = '2-m temperature'
unit     = r"${\rm ^\circ C}$"     # r + symbol of the unit of the values ("m", "%", "${\rm ^\circ C}$", "m/s", "${\rm W/m^2}$", ...)
val_min  =  -4
val_max  =  18
val_int  =   2
facteur_add = -273.15
# Output image information
image_output_file = "./image.gif"
image_output_dpi = 150
# ================================================

netcdf_file = netCDF4.Dataset(filename)

# Read RP projection attributes
proj_field = netcdf_file.variables['rotated_pole']
grid_north_pole_longitude = proj_field.grid_north_pole_longitude
grid_north_pole_latitude = proj_field.grid_north_pole_latitude
north_pole_grid_longitude = proj_field.north_pole_grid_longitude


# Read fields
main_field = netcdf_file.variables[varname][:,:]
lat_d = netcdf_file.variables['lat'][:]
lon_d = netcdf_file.variables['lon'][:]

netcdf_file.close()                                     # Close file

# 2-D Mapping - if needed
# -----------------------
m = ccrs.RotatedPole(pole_longitude=grid_north_pole_longitude,pole_latitude=grid_north_pole_latitude)

# Figure settings - if needed
# ---------------------------
figsize = (5, 4.4)      # Figure size
fig = plt.figure(figsize=figsize)

# Set projection defined by the cartopy object
ax = plt.axes(projection=m)

# Plotting - if needed
# --------------------

# Set corners of the maps
xll, yll = m.transform_point(lon_d[0, 0],lat_d[0, 0],ccrs.PlateCarree())
xur, yur = m.transform_point(lon_d[-1, -1],lat_d[-1, -1], ccrs.PlateCarree())

# Set geographic features
ax.add_feature(cfeature.OCEAN.with_scale('50m'))      # couche ocean
ax.add_feature(cfeature.LAND.with_scale('50m'))       # couche land
ax.add_feature(cfeature.LAKES.with_scale('50m'),edgecolor='black',facecolor='none')      # couche lac    
ax.add_feature(cfeature.BORDERS.with_scale('50m'))    # couche frontieres
#ax.add_feature(cfeature.RIVERS.with_scale('50m'))     # couche rivières 
coast = cfeature.NaturalEarthFeature(category='physical', scale='10m',facecolor='none', name='coastline') # Couche côtières
ax.add_feature(coast, edgecolor='black')

states_provinces = cfeature.NaturalEarthFeature(category='cultural',name='admin_1_states_provinces_lines',scale='10m',facecolor='none') # Couche provinces
ax.add_feature(states_provinces, edgecolor='black')

ax.set_extent([xll, xur, yll, yur], crs=m)

# Set colors
# ----------
# Set color map
cmap = plt.cm.jet

# Set min, max and interval for color plot 
ints = ( (val_max-val_min+val_int) / val_int)
color_bnds  = np.linspace(val_min, val_max, int(ints))
color_ticks = np.linspace(val_min, val_max, int(ints))

norm = colors.BoundaryNorm(boundaries=color_bnds, ncolors=256)


# Plot figure
cs = plt.pcolormesh(lon_d, lat_d, main_field+facteur_add,norm=norm,cmap=cmap,transform=ccrs.PlateCarree())

# Plot colorbar
cb = plt.colorbar(mappable=cs,ticks=color_ticks,orientation="vertical",pad=0.05,shrink=0.8, extend="both")
cb.ax.set_title(unit)

# Plot title
plt.title(title, fontsize=12)

# To help the layout of the figure after saving
fig.canvas.draw()
plt.tight_layout()  # To help with the layout of the figure after saving

# Save figure
fig.savefig('python.png', dpi=image_output_dpi, format='png') # Most backends support png, pdf, ps, eps and svg
os.system ('convert python.png ' + image_output_file)         # Convert python.png to 'image_output_file', i.e. .gif