#!/usr/bin/env python
import os,sys
import xarray
import numpy
from matplotlib import pylab
import argparse
try:
    import seaborn as sns
    seaborn_available = True
    sns.set()
    sns.set_style("darkgrid")
#   sns.set_style("white")
except ImportError:
    seaborn_available = False
# OPTIONS ======================================================================
parser = argparse.ArgumentParser()
parser.add_argument("-i","--input",    help="input filename",required=True)
parser.add_argument("-v","--varname",  help="varname to plot from the input file",required=True)
parser.add_argument("-t","--timestep", type=int, help="timstep to plot", default=0)
parser.add_argument("-o","--output",   type=str, help="output image filename", default=None)
parser.add_argument("-m","--minmax",   type=int, help="min/max value for the colorbar", nargs=2 , default=[])
parser.add_argument("-l","--nlev",     type=int, help="number of colors", default=11)
parser.add_argument("-d","--debug",    action="store_true", help="get debugging output", default=False)
# MAIN SCRIPT ==================================================================
options = parser.parse_args()

if options.debug:
    print(options)
    print(sys.argv)

fig = pylab.figure()
if os.path.splitext(options.input)[1] in ['.grb','.grib','.grb1','.grib1','.grb2','.grib2']:
    data = xarray.open_dataset(options.input,engine='cfgrib')
else:
    data = xarray.open_dataset(options.input)

if options.debug:
    print(data)

if options.output != None:
    import matplotlib
    matplotlib.use('Agg')

if [] != options.minmax:
    levels = numpy.linspace(options.minmax[0],options.minmax[1],options.nlev)
    data[options.varname][options.timestep,:,:].plot(levels=levels)
else:
    data[options.varname][options.timestep,:,:].plot()

if options.output != None:
    pylab.savefig(options.output)
else:
    pylab.show()

pylab.close(fig)