#!/bin/env/python

import grib2io
import pyproj
import cartopy.crs as ccrs
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import cartopy.feature as cfeature
import matplotlib
#matplotlib.use('Agg')
import io
import matplotlib.pyplot as plt
from PIL import Image
import matplotlib.image as image
from matplotlib.gridspec import GridSpec
import numpy as np
import time,os,sys,multiprocessing
import multiprocessing.pool
import ncepy
from scipy import ndimage
from netCDF4 import Dataset
import cartopy
import rrfs_plot_utils

#--------------Set some classes------------------------#
# Make Python process pools non-daemonic
class NoDaemonProcess(multiprocessing.Process):
  # make 'daemon' attribute always return False
  @property
  def daemon(self):
    return False

  @daemon.setter
  def daemon(self, value):
    pass

class NoDaemonContext(type(multiprocessing.get_context())):
  Process = NoDaemonProcess

# We sub-class multiprocessing.pool.Pool instead of multiprocessing.Pool
# because the latter is only a wrapper function, not a proper class.
class MyPool(multiprocessing.pool.Pool):
  def __init__(self, *args, **kwargs):
    kwargs['context'] = NoDaemonContext()
    super(MyPool, self).__init__(*args, **kwargs)


#--------------Define some functions ------------------#

def clear_plotables(ax,keep_ax_lst,fig):
  #### - step to clear off old plottables but leave the map info - ####
  if len(keep_ax_lst) == 0 :
    print("clear_plotables WARNING keep_ax_lst has length 0. Clearing ALL plottables including map info!")
  cur_ax_children = ax.get_children()[:]
  if len(cur_ax_children) > 0:
    for a in cur_ax_children:
      if a not in keep_ax_lst:
       # if the artist isn't part of the initial set up, remove it
        a.remove()

def compress_and_save(filename):
  #### - compress and save the image - ####
#  ram = io.StringIO()
  ram = io.BytesIO()
  plt.savefig(ram, format='png', bbox_inches='tight', dpi=300)
#  plt.savefig(filename, format='png', bbox_inches='tight', dpi=300)
  ram.seek(0)
  im = Image.open(ram)
  im2 = im.convert('RGB').convert('P', palette=Image.ADAPTIVE)
  im2.save(filename, format='PNG')

#-------------------------------------------------------#

# Necessary to generate figs when not running an Xserver (e.g. via PBS)
plt.switch_backend('agg')

# Read date/time and forecast hour from command line
ymdh = str(sys.argv[1])
ymd = ymdh[0:8]
year = int(ymdh[0:4])
month = int(ymdh[4:6])
day = int(ymdh[6:8])
hour = int(ymdh[8:10])
cyc = str(hour).zfill(2)
print(year, month, day, hour)

fhr = int(sys.argv[2])
fhrm1 = fhr - 1
fhrm2 = fhr - 2
fhrm6 = fhr - 6
fhrm24 = fhr - 24
fhour = str(fhr).zfill(2)
fhour1 = str(fhrm1).zfill(2)
fhour2 = str(fhrm2).zfill(2)
fhour6 = str(fhrm6).zfill(2)
fhour24 = str(fhrm24).zfill(2)
print('fhour '+fhour)

# Define the input files
data1 = grib2io.open('/lfs/h2/emc/ptmp/Benjamin.Blake/work/test1/gfs.t00z.pres.f0'+fhour+'.grib2')
data2 = grib2io.open('/lfs/h2/emc/ptmp/Benjamin.Blake/work/test2/gfs.t00z.pres.f0'+fhour+'.grib2')

msg = data2.select(shortName='HGT', level='0C isotherm')[0] 	# msg is a Grib2Message object
lat2,lon2 = msg.latlons()

# Forecast valid date/time
itime = ymdh
vtime = rrfs_plot_utils.ndate(itime,int(fhr))

# Specify plotting domains
#domains = ['conus','BN','CE','CO','LA','MA','NC','NE','NW','OV','SC','SE','SF','SP','SW','UM']
domains=['globe']

###################################################
# Read in all variables and calculate differences #
###################################################
t1a = time.perf_counter()

# Surface height
hgtsfc_1 = data1.select(shortName='HGT',level='surface')[0].data
hgtsfc_2 = data2.select(shortName='HGT',level='surface')[0].data

# Lowest freezing level
frzlow_1 = data1.select(shortName='HGT',level='0C isotherm')[0].data
frzlow_2 = data2.select(shortName='HGT',level='0C isotherm')[0].data

# Highest freezing level
frzhi_1 = data1.select(shortName='HGT',level='highest tropospheric freezing level')[0].data
frzhi_2 = data2.select(shortName='HGT',level='highest tropospheric freezing level')[0].data

hgtdif = frzhi_2 - frzlow_2


t2a = time.perf_counter()
t3a = round(t2a-t1a, 3)
print(("%.3f seconds to read all messages") % t3a)

# colors for difference plots, only need to define once
difcolors = ['blue','#1874CD','dodgerblue','deepskyblue','turquoise','white','white','#EEEE00','#EEC900','darkorange','orangered','red']
difcolors2 = ['white']
difcolors3 = ['blue','dodgerblue','turquoise','white','white','#EEEE00','darkorange','red']

########################################
#    START PLOTTING FOR EACH DOMAIN    #
########################################

def main():

  # Number of processes must coincide with the number of domains to plot
#  pool = multiprocessing.Pool(len(domains))
#  pool = MyPool(len(domains))
#  pool.map(plot_all,domains)
  p = MyPool(len(domains))
  p.map(plot_all,domains)

def plot_all(domain):

  global dom
  dom = domain
  print(('Working on '+dom))

  global fig,axes,ax1,keep_ax_lst_1,xextent,yextent,im,par,transform
  fig,axes,ax1,keep_ax_lst_1,xextent,yextent,im,par,transform = create_figure()

  # Split plots into 2 sets with multiprocessing
  sets = [1]
  pool2 = multiprocessing.Pool(len(sets))
  pool2.map(plot_sets,sets)

def create_figure():

  # Map corners for each domain
  if dom == 'conus':
    llcrnrlon = -125.5
    llcrnrlat = 20.0 
    urcrnrlon = -63.5
    urcrnrlat = 51.0
    cen_lat = 35.4
    cen_lon = -97.6
    xextent=-2200000
    yextent=-675000
  elif dom == 'namerica': 
    llcrnrlon = -160.0
    llcrnrlat = 15.0 
    urcrnrlon = -55.0
    urcrnrlat = 65.0
    cen_lat = 35.4
    cen_lon = -105.0
    xextent = -3700000
    yextent = -2500000
    offset = 1
  elif dom == 'gsl':
    llcrnrlon = -115.0
    llcrnrlat = 38.0 
    urcrnrlon = -110.0
    urcrnrlat = 42.5
    cen_lat = 40.0
    cen_lon = -112.5
    xextent=-2200000
    yextent=-675000
  elif dom == 'globe':
    llcrnrlon = -179.9
    llcrnrlat = -90.0
    urcrnrlon = 180.0
    urcrnrlat = 90.0
    cen_lat = 0.0
    cen_lon = 0.0
    xextent=-10000
    yextent=-10000

  # create figure and axes instances
  fig = plt.figure(figsize=(4,4))
  gs = GridSpec(4,4,wspace=0.0,hspace=0.0)
  im = image.imread('/lfs/h2/emc/lam/noscrub/Benjamin.Blake/python.rrfs/noaa.png')
  par = 1

  # Define where Cartopy maps are located
  cartopy.config['data_dir'] = '/lfs/h2/emc/lam/noscrub/Benjamin.Blake/python/NaturalEarth'

  back_res='50m'
  back_img='off'

  # set up the map background with cartopy
  if dom == 'conus':
    extent = [llcrnrlon-1,urcrnrlon-6,llcrnrlat,urcrnrlat+1]
  elif dom == 'namerica':
    extent = [-176.,0.,0.5,45.]
  else:
    extent = [llcrnrlon,urcrnrlon,llcrnrlat,urcrnrlat]

#  myproj = ccrs.Orthographic(central_longitude=-114, central_latitude=54.0, globe=None)

#  myproj = ccrs.LambertConformal(central_longitude=cen_lon,central_latitude=cen_lat,
#           false_easting=0.0,false_northing=0.0,globe=None)
  myproj = ccrs.PlateCarree(central_longitude=cen_lon, globe=None)

  ax1 = fig.add_subplot(gs[0:4,0:4], projection=myproj)
  ax1.set_extent(extent)
  axes = [ax1]

  fline_wd = 0.5  # line width
  fline_wd_lakes = 0.15  # line width
  falpha = 0.5    # transparency

  # natural_earth
  lakes=cfeature.NaturalEarthFeature('physical','lakes',back_res,
                    edgecolor='black',facecolor='none',
                    linewidth=fline_wd_lakes)
  coastlines=cfeature.NaturalEarthFeature('physical','coastline',
                    back_res,edgecolor='black',facecolor='none',
                    linewidth=fline_wd,alpha=falpha)
  states=cfeature.NaturalEarthFeature('cultural','admin_1_states_provinces',
                    back_res,edgecolor='black',facecolor='none',
                    linewidth=fline_wd,alpha=falpha)

  # All lat lons are earth relative, so setup the associated projection correct for that data
#  transform = ccrs.RotatedPole(pole_longitude=67.0, pole_latitude=35.0)
  transform = ccrs.PlateCarree()

  # high-resolution background images
  if back_img=='on':
    img = plt.imread('/lfs/h2/emc/lam/noscrub/Benjamin.Blake/python/NaturalEarth/raster_files/NE1_50M_SR_W.tif')
    ax1.imshow(img, origin='upper', transform=transform)

  ax1.add_feature(cfeature.LAND, linewidth=0, facecolor='white')
  ax1.add_feature(cfeature.OCEAN, linewidth=0, facecolor='lightgray')
  ax1.add_feature(cfeature.LAKES, edgecolor='black', linewidth=fline_wd_lakes, facecolor='lightgray',zorder=0)
  ax1.add_feature(lakes)
  ax1.add_feature(states)
  ax1.add_feature(coastlines)

  # Map/figure has been set up here, save axes instances for use again later
  keep_ax_lst_1 = ax1.get_children()[:]


  return fig,axes,ax1,keep_ax_lst_1,xextent,yextent,im,par,transform


def plot_sets(set):
# Add print to see if dom is being passed in
  print(('plot_sets dom variable '+dom))

  global fig,axes,ax1,keep_ax_lst_1,xextent,yextent,im,par,transform

  if set == 1:
    plot_set_1()
  elif set == 2:
    plot_set_2()
  elif set == 3:
    plot_set_3()

def plot_set_1():
  global fig,axes,ax1,keep_ax_lst_1,xextent,yextent,im,par,transform


#################################
  # Plot freezing level
#################################
  t1dom = time.perf_counter()
  t1 = time.perf_counter()

  print(('Working on freezing level for '+dom))

  units = 'gpm'
# Use this pair for frzhi - frzlow
  clevs = [0,10,50,100,250,500,750,1000,1500,2000,2500,3000]
  colorlist = ['white','blue','dodgerblue','cyan','mediumspringgreen','#EEEE00','#EEC900','darkorange','crimson','darkred','darkviolet']
# Use this pair for frzlow - hgtsfc or frzhi - hgtsfc
#  clevs = [-4000,-3000,-2000,-1000,-500,0,500,1000,2000,3000,4000,5000,6000]
#  colorlist = ['darkblue','blue','dodgerblue','cyan','mediumspringgreen','#FAFAD2','#EEEE00','#EEC900','darkorange','crimson','darkred','darkviolet']


  cm = matplotlib.colors.ListedColormap(colorlist)
  norm = matplotlib.colors.BoundaryNorm(clevs, cm.N)

  xmin, xmax = ax1.get_xlim()
  ymin, ymax = ax1.get_ylim()
  xmax = int(round(xmax))
  ymax = int(round(ymax))

  cs_1 = ax1.pcolormesh(lon2,lat2,hgtdif,transform=transform,cmap=cm,norm=norm)
  cs_1.cmap.set_under('white',alpha=0.)
  cs_1.cmap.set_over('black')
  cbar1 = fig.colorbar(cs_1,ax=ax1,orientation='horizontal',pad=0.01,shrink=0.9,ticks=clevs,extend='max')
  cbar1.set_label(units,fontsize=6)
  cbar1.ax.tick_params(labelsize=5)
  ax1.text(.5,1.03,'Highest Freezing Level - Lowest Freezing Level (new) \n initialized: '+itime+' valid: '+vtime + ' (f'+fhour+')',horizontalalignment='center',fontsize=6,transform=ax1.transAxes,bbox=dict(facecolor='white',alpha=0.85,boxstyle='square,pad=0.2'))
#  ax1.text(.5,0.03,'Experimental Product - Not Official Guidance',horizontalalignment='center',fontsize=6,color='red',transform=ax1.transAxes,bbox=dict(facecolor='white',color='white',alpha=0.85,boxstyle='square,pad=0.2'))
  ax1.imshow(im,aspect='equal',alpha=0.5,origin='upper',extent=(xmin,xextent,ymin,yextent),zorder=4)


  rrfs_plot_utils.convert_and_save('frzlvl_'+dom+'_f'+fhour)
  t2 = time.perf_counter()
  t3 = round(t2-t1, 3)
  print(('%.3f seconds to plot lowest freezing level for: '+dom) % t3)



  t3dom = round(t2-t1dom, 3)
  print(("%.3f seconds to plot all set 1 variables for: "+dom) % t3dom)
  plt.clf()


################################################################################

main()