#! /usr/bin/env python
# sources:
# http://matplotlib.org/basemap/users/examples.html
# http://matplotlib.org/users/colormapnorms.html
from mpl_toolkits.basemap import Basemap
import numpy as np
import math as m
import matplotlib.pyplot as plt
#import matplotlib.dates as dates
#import datetime
#from matplotlib.dates import date2num
from matplotlib import ticker
#import pandas
from pylab import *
import netCDF4
import sys, getopt

def main(argv):
  inputfile = ''
  outputdir = ''
  errmsg='Use: '+str(sys.argv[0])+' -i <inputfile> -o <outputdir>'
  try:
     opts, args = getopt.getopt(argv,"h:i:o:")
  except getopt.GetoptError:
     print(errmsg)
     sys.exit(2)
  for opt, arg in opts:
     if opt == '-h':
        print(errmsg)
        sys.exit()
     elif opt in ("-i"):
        inputfile = arg
     elif opt in ("-o"):
        outputdir = arg
  if len(inputfile) == 0 or len(outputdir) == 0:
    print('Please specify an input file and an output directory.')
    print(errmsg)
    sys.exit()
  print('Input file is ', inputfile)
  print('Output directory is ', outputdir)

  # MAIN PROGRAM
  # -----------------------------------------------------------------
  
  ncfile=inputfile
  shortname=(ncfile.rsplit('/',1))[1].rsplit('_',-1)[0].replace(".","-")
  
  # LOADING THE GCM RESULTS
  # -----------------------------------------------------------------
  
  ncgcm = netCDF4.Dataset(ncfile)
  longi0 = ncgcm.variables['lon'][:] 
  longi = longi0
  lati0 = ncgcm.variables['lat'][:]
  lati = lati0
  snowgcm = ncgcm.variables['snow'][:] # precip rate in mm/hr
  
  ncpath='/thredds/ipsl/fabric/lmdz/AXE4/'
  ncfractlic=ncpath + 'fract_lic_144x142.nc'
  ncmask = netCDF4.Dataset(ncfractlic)
  fract_lic = ncmask.variables['fract_lic'][:]
  fract_lic = fract_lic[0,:,:] # constant anyway
  
  # DISPLAYING THE RESULTS
  # -----------------------------------------------------------------
  
  # POLAR MAP
  # source: http://polar.ncep.noaa.gov/global/examples/usingpython.shtml 
  #         http://matplotlib.org/basemap/users/pstere.html 
  # -----------------------------------------------------------------
  
  # Some postprocessing
  #>>> snow.shape
  #(1, 143, 144)
  # Convert missing values to nan
  snowgcm = np.nanmean(snowgcm,axis=0)*86400.*365.
  snowgcm[fract_lic==0]=np.nan
  data = snowgcm
  
  # Setup north polar stereographic basemap.
  # The longitude lon_0 is at 6-o'clock, and the
  # latitude circle boundinglat is tangent to the edge
  # of the map at lon_0. Default value of lat_ts
  # (latitude of true scale) is pole.
  m = Basemap(projection='spstere',boundinglat=-60, \
    lon_0=180.,resolution='c')
  
  # Convert the lat/lon values to x/y projections.
  
  Lon, Lat = meshgrid(longi,lati)
  x, y = m(Lon,Lat)
  
  # draw filled contours.
  clevs = [0.,20,50,100,200,300,400,500,700,1E3,1500,2E3,7E3]
  #clevs = [0.,0.25,0.5,0.75,1.,1.25,1.5,2.,3.,5.,10.,20]
  cs = m.contourf(x,y,data,clevs, \
  # norm = mpl.colors.PowerNorm(gamma=0.15), \
  # norm = mpl.colors.LogNorm(vmin=clevs[0],vmax=clevs[len(clevs)-1]), \
  # norm = mpl.colors.Normalize(vmin=clevs[0],vmax=clevs[len(clevs)-1]), \
   norm = mpl.colors.BoundaryNorm(boundaries=clevs, ncolors=256), \
   cmap='jet')
  #cs = m.pcolor(x,y,data, \
  #  norm = mpl.colors.Normalize(vmin=clevs[0],vmax=clevs[len(clevs)-1]))
  #cs = m.contourf(x,y,data,clevs)
  #cs = m.contour(x,y,data,clevs,colors='k',linewidths=1.)
  # Add a colorbar and title, and number of bins
  cbar = m.colorbar(cs,location='bottom',pad="5%")
  #cbar.set_label('mm')
  #cbar.locator = ticker.MaxNLocator(nbins=15)
  #cbar.update_ticks()
  
  # Next, plot the field using the fast pcolormesh routine and set the colormap to jet.
  #cs = m.pcolormesh(x,y,data,shading='flat', \
  #cmap=plt.cm.jet, \
  #vmin=0., vmax=100)
  
  # Add a coastline and axis values.
  m.drawcoastlines()
  #m.fillcontinents()
  m.drawmapboundary()
  m.drawparallels(np.arange(-90.,95.,5.), \
  labels=[1,0,0,0])
  m.drawmeridians(np.arange(-180.,180.,30.), \
  labels=[0,0,0,1])
  
  plt.title('Precip in mm/yr ('+shortname+')')
  
  savefig(outputdir+'/'+'precipSH-'+shortname+'.png', bbox_inches='tight')
  #plt.show()
  
  # -----------------------------------------------------------------
  # -----------------------------------------------------------------
  # -----------------------------------------------------------------

# MAIN PROGRAM
# Just calling the main function
if __name__ == "__main__":
   main(sys.argv[1:])