#!/usr/bin/env python3
#
# Before opening Python you need to first load the Python version (in this case Python3) you want to use with:
#   module load python3/miniconda3
# and then load the module needed to read RPN files with Python with:
#   module load python3/python-rpn
#   module load python3/outils-divers
# and to get access to basic Python packages source:
#   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 rpnpy.librmn.all as rmn             # Module to read RPN files
from rotated_lat_lon import RotatedLatLon  # Module to project field on native grid (created by Sasha Huziy)

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

# Example parameters which need to get adjusted !!!
# ================================================
filename = '/sca/compilers_and_tools/python/demos-python3/data/dmCAN_Rockies2_3km_P3_2013_ERAInt_201305_moyenne'    # Name of RPN file to read
varname  = 'TT'    # 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
# Output image information
image_output_file = "./test.gif"
image_output_dpi = 150
# ================================================



# Read one record
# ---------------
fid = rmn.fstopenall(filename,rmn.FST_RO)   # Open the file
rec = rmn.fstlir(fid,nomvar=varname)        # Read the full record of variable 'varname'
field = rec['d']                            # Assign 'field' to the data of 'varname'


# Read 2-D latitudes & longitudes - if needed
mygrid = rmn.readGrid(fid,rec)              # Get the grid information for the (LAM) Grid -- Reads the tictac's
latlondict = rmn.gdll(mygrid)               # Create 2-D lat and lon fields from the grid information
lat = latlondict['lat']                     # Assign 'lat' to 2-D latitude field
lon = latlondict['lon']                     # Assign 'lon' to 2-D longitude field

# Get grid rotation for projection of 2-D field for mapping - if needed
tics = rmn.fstlir(fid,nomvar='^^', ip1=rec['ig1'],ip2=rec['ig2'],ip3=rec['ig3']) # Read corresponding tictac's
 
# Close RPN input file
rmn.fstcloseall(fid)                        # Close the RPN file


# 2-D Mapping - if needed
# -----------------------
# Get positions of rotated equator from IG1-4 of the tictac's
(Grd_xlat1,Grd_xlon1,Grd_xlat2,Grd_xlon2) = rmn.cigaxg('E', tics['ig1'],tics['ig2'],tics['ig3'],tics['ig4'])

# Use Sasha's RotatedLatLon to get the rotation matrix
rll = RotatedLatLon( lon1=Grd_xlon1, lat1=Grd_xlat1, lon2=Grd_xlon2, lat2=Grd_xlat2) # the params come from gemclim_settings.nml
# Use Sasha's get_cartopy_projection_obj to get the cartopy object for the projection and domain defined by the coordinates
m = rll.get_cartopy_projection_obj()

# 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[0, 0],lat[0, 0],ccrs.PlateCarree())
xur, yur = m.transform_point(lon[-1, -1],lat[-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'))      # 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, lat, field,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