#! /usr/bin/python

#This code loads default plotting maps in cartopy for convenience. 20210205-0224. MJE.

import numpy as np
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
from scipy import interpolate
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER

#Simple definition to load all of the basic map properites. Inputs can be a 
#grid of latitude/longitude coordinates (2-D) or lon/lat minimum and maximum.
#MJE. 20210205-0309.
#
#####################INPUT FILES FOR plot_map################################
# lat   = latitude data (2-D grid or [min,max])
# lon   = longitude data (2-D grid or [min,max])
#####################OUTPUT FILES FOR plot_map################################
# fig   = figure properties
# ax    = axes properties
##################################################################################

def plot_map(lat,lon):

    lat = np.array(lat)
    lon = np.array(lon)

    fig = plt.figure(figsize=(12,6)) #Create Figure
    ax  = plt.axes(projection=ccrs.Mercator()) #Create geoaxes to the projection

    #Draw coastlines, state and country boundaries, edge of map.
    ax.coastlines(edgecolor='grey',linewidth=0.5)
    ax.add_feature(cf.STATES,edgecolor='grey',linewidth=0.3)
    ax.add_feature(cf.BORDERS,edgecolor='grey',linewidth=0.3)
    #Determine resolution of gridlines
    if len(lon.shape) == 2:
        gridsize = np.ceil((np.ceil(np.nanmax(lon))-np.floor(np.nanmin(lon))))
    else:
        gridsize = np.ceil((np.ceil(lon[1])-np.floor(lon[0])))

    if gridsize >= 250:
        gridline_res = 40
    elif gridsize < 250 and gridsize >= 200:
        gridline_res = 30
    elif gridsize < 200 and gridsize >= 150:
        gridline_res = 20
    elif gridsize < 150 and gridsize >= 100:
        gridline_res = 10
    elif gridsize < 100:
        gridline_res = 5
    if len(lon.shape) == 2:
        xticks = np.arange(np.nanmin(lon)-np.nanmin(lon)%gridline_res,np.nanmax(lon)+np.nanmax(lon)%gridline_res+gridline_res,gridline_res)
        yticks = np.arange(np.nanmin(lat)-np.nanmin(lat)%gridline_res,np.nanmax(lat)+np.nanmax(lat)%gridline_res+gridline_res,gridline_res)
    else:
        xticks = np.arange(lon[0]-lon[0]%gridline_res,lon[1]+lon[1]%gridline_res+gridline_res,gridline_res)
        yticks = np.arange(lat[0]-lat[0]%gridline_res,lat[1]+lat[1]%gridline_res+gridline_res,gridline_res)

    gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                  linewidth=0.5, color='black', alpha=0.5, linestyle='--',xlocs=xticks,ylocs=yticks)
    gl.top_labels = False
    gl.right_labels = False
    gl.bottom_labels = True
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    gl.xlabel_style = {'size': 12}
    gl.ylabel_style = {'size': 12}
    if len(lon.shape) == 2:
        ax.set_extent([np.nanmin(lon), np.nanmax(lon), np.nanmin(lat), np.nanmax(lat)], crs=ccrs.PlateCarree())
    else:
        ax.set_extent([lon[0], lon[1], lat[0], lat[1]], crs=ccrs.PlateCarree())

    return(fig,ax)