#! /usr/bin/python

#Simple code to make 2D plot. This code has been upgraded to be compatable with python3.
#MJE. Created 20200511 and upgraded 20210208,20210224.

import copy
import subprocess
import os
import numpy as np
import math
import cartopy_maps
import cartopy.crs as ccrs                   # import projections
import cartopy.feature as cf                 # import features
import matplotlib as mpl
mpl.use('Agg') #So plots can be saved in cron
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from matplotlib.path import Path
from matplotlib.cm import get_cmap
from scipy.io import netcdf
from netCDF4 import Dataset

##################################################################################
#Need a simple 2D plotter that creates the (netcdf) plot on a file and 
#throws in to the web.
#MJE. 20200511.
#
#Update: Converted to python 3 and added ability to load lat/lon directly from 
#netCDF file. 20210224. MJE.
#
#####################INPUT FILES FOR readCONUSmask################################
# latlon_dims   = latitude longitude plotting dimensions
# file_open     = netcdf file, full directory
# file_var      = 2-D variable file name inside netcdf file
# fig_put       = full directory to scp to web
# fig_name      = figure name
# fig_colorbar  = colorbar in increments to be plotted
##################################################################################

def Make2DPlot(lat,lon,file_open,file_var,fig_put,fig_name,fig_colorbar):

    #Load the dataset
    f         = Dataset(file_open, "a", format="NETCDF4")

    #If latitude/longitude data is not specified at input, try loading it from the netCDF file
    if len(lat) == 0:
        lat          = f.variables['lat'][:]
        lon          = f.variables['lon'][:]
        lon          = np.flipud(np.tile(lon,(lat.shape[0],1)))
        lat          = np.flipud(np.transpose(np.tile(lat,(lon.shape[1],1))))

    temp      = (np.array(f.variables[file_var][:]))*1

    #Create colormap
    cmap = copy.copy(get_cmap("rainbow"))
    cmap.set_under('white')

    #Definition to create default map
    fig = plt.figure(figsize=(12,6)) #Create Figure
    ax  = plt.axes(projection=ccrs.Mercator(),alpha=1) #Create geoaxes to the projection
    cartopy_maps.plot_map(lat,lon,ax) 

    #Create the contour fill and colorbar
    ax_c = plt.contourf(lon, lat, temp, levels=fig_colorbar, extend='min',cmap=cmap,transform=ccrs.PlateCarree())
    plt.colorbar(ax_c,ax=ax)
    ax.set_extent([np.nanmin(lon), np.nanmax(lon), np.nanmin(lat), np.nanmax(lat)], crs=ccrs.PlateCarree())
    plt.savefig('/export/hpc-lw-dtbdev5/merickson/Desktop/'+fig_name, bbox_inches='tight',dpi=200)
    subprocess.call('scp /export/hpc-lw-dtbdev5/merickson/Desktop/'+fig_name+' '+fig_put, \
        stdout=open(os.devnull, 'wb'),stderr=open(os.devnull, 'wb'), shell=True)
    plt.close()
    subprocess.call('rm -rf /export/hpc-lw-dtbdev5/merickson/Desktop/'+fig_name, shell=True)