#---------------------------------------------------------------
# This script plots the first time/level of all 3D variables on 
# the given file using PyNGL.
#---------------------------------------------------------------
import numpy as np
import Nio, Ngl, os
import create_html

def create_plots():
#    html_name = 'all_plots_grib1_ngl'
#    plots_dir = 'plots_grib1_ngl'
#    filename  = '../Data/fnl_20070703_06_00.grib1'

    html_name = 'all_plots_gdas1_fnl_ngl'
    plots_dir = 'plots_gdas1_fnl_ngl'
    filename  = '../Data/gdas1.fnl0p25.2016030100.f00.grib2'

    if not os.path.exists(plots_dir):
        print("Create directory '%s' before running this script." % plots_dir)
        return 1

    if not os.path.exists(filename):
        print("File '%s' doesn't exist." % filename)
        return 1

#---Try to open the file
    try:
        f = Nio.open_file(filename)
    except:
        print("Cannot open file '%s'." % filename)
        return 1

#    print(f.variables.keys())

#---Now start reading variables and creating plots.
    lat = f.variables["lat_0"][:]
    lon = f.variables["lon_0"][:]

#---Set plot resources common to all plots
    res          = Ngl.Resources()

    res.sfYArray = lat
    res.sfXArray = lon

#---Contour resources
    res.cnFillOn             = True
    res.cnFillPalette        = "MPL_jet"          # change color map
    res.cnLinesOn            = False
    res.cnLineLabelsOn       = False
    res.cnFillMode           = "RasterFill"       # for faster plotting

#---Labelbar and title resources
    res.lbOrientation        = "horizontal"
    res.lbLabelFontHeightF   = 0.01
    res.tiMainFontHeightF    = 0.015

#---Map resources
    res.mpGridAndLimbOn        = False
    res.mpCenterLonF           = 180

#---Loop through each variable and create a plot
    for varname in f.variables.keys():
        var = f.variables[varname]

        if (len(var.shape) == 2):
            title = "%s: %s (%s)" % (varname,var.long_name,var.units)
            print "plotting %s" % title

#---Open PNG, set title, create plot
            wks = Ngl.open_wks("png",'%s/%s' % (plots_dir,varname))
            res.tiMainString      = title
            map = Ngl.contour_map(wks,var,res)
    
            Ngl.destroy(wks)    # close plot

#---Create HTML file to display images.
    html = create_html.create_html(html_name,'Plots created with PyNGL',plots_dir)

#---Return from create_plots
    return 0

if __name__ == '__main__':
    create_plots()