#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jan  3 15:51:39 2018
Modified on Jul 27, 2022 to add command line filename and outputdir

@author: jamiller, SGT Inc. for NASA/GSFC/Code 615

usage: python ncdftest.py -i [filename] -o [outputdir]
    output dir is optional, otherwise files will write to current dir
"""

import sys, getopt # routines that let you pass command line parameters
from netCDF4 import Dataset # netcdf routines that you may need to install into python
from pathlib import Path # routines that allow you to check for directories and files

def main(argv):
   inputfile = ''
   outputdir = ''
   try:
      opts, args = getopt.getopt(argv,"hi:o:",["ifile=","odir="])
   except getopt.GetoptError:
      print('test.py -i <inputfile> -o <outputdir>')
      sys.exit(2)
   for opt, arg in opts:
      if opt == '-h':
         print('test.py -i <inputfile> -o <outputdir>')
         sys.exit()
      elif opt in ("-i", "--ifile"):
         inputfile = arg
      elif opt in ("-o", "--odir"):
         outputdir = arg
   rootgrp = Dataset(inputfile, "r") # open file for reading
   print(rootgrp.variables.keys()) # show the keys / variables in the file
   print(rootgrp.variables['Melt']) # show the information on the 'melt' variable

   for attr in rootgrp.ncattrs(): # show the global attributes in the file
       print(attr, '=', getattr(rootgrp, attr))

   year = getattr(rootgrp, 'Year') # pull year for plots
   vers = getattr(rootgrp, 'Melt Version') # pull version for plot savefiles

   # extract variables
   melt = rootgrp.variables['Melt'][:]
   earlymelt = rootgrp.variables['Earlymelt'][:]
   freeze = rootgrp.variables['Freeze'][:]
   earlyfreeze = rootgrp.variables['Earlyfreeze'][:]

   rootgrp.close() # close out the netcdf file
   
   #if outputdir is named, does it exist? Create it if not
   Path(outputdir).mkdir(parents=True, exist_ok=True)

   import matplotlib.pyplot as plt # plot libraries
   imgplot = plt.imshow(melt, clim=(75, 210)) # do an image plot scaled from 75 to 210
   plt.colorbar() # give it a colorbar
   plt.title(year+ ' Melt') # title it
   plt.savefig(outputdir+year+vers+'meltexample.pdf', dpi=300) # save it
   plt.show() # show it and clear the graphics buffer
   # repeat for the other variables in the melt file
   imgplot2 = plt.imshow(earlymelt, clim=(75, 210))
   plt.colorbar()
   plt.title(year+ ' Early Melt')
   plt.savefig(outputdir+year+vers+'earlymeltexample.pdf', dpi=300)
   plt.show()
    
   imgplot3 = plt.imshow(freeze, clim=(210, 410))
   plt.colorbar()
   plt.title(year+ ' Freeze')
   plt.savefig(outputdir+year+vers+'freezeexample.pdf', dpi=300)
   plt.show()

   imgplot4 = plt.imshow(earlyfreeze, clim=(210, 410))
   plt.colorbar()
   plt.title(year+ ' Early Freeze')
   plt.savefig(outputdir+year+vers+'earlyfreezeexample.pdf', dpi=300)
   plt.show()


if __name__ == "__main__":
   main(sys.argv[1:])