#----------------------------------------------------------------------
# Read and plots all 2D variables on a GRIB2 file using PyNIO/matplotlib.
#
# Curvilinear grid, 880 x 1760
#
# This script is expecting the lat/lon arrays on the file to be 
# 2D (curvilinear). If they are not, then see the "meshgrid" 
# line below, which is commented out.
#----------------------------------------------------------------------
import numpy as np, sys, os
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import Nio
import create_html

# Initialize variables
html_name = 'all_plots_grib2_mpl'
plots_dir = 'plots_grib2_mpl'
if not os.path.exists(plots_dir):
  print("Create directory '%s' before running this script." % plots_dir)
  sys.exit()

# Read data
dir      = "../Data/"
filename = "gdas1.t18z.sfluxgrbf03.grib2"
a   = Nio.open_file("%s/%s" %(dir,filename))
# You may need to change these two lines if your lat/lon variables
# have different names on the file.
lat = a.variables["gridlat_236"][:]
lon = a.variables["gridlon_236"][:]
#
# If your lat/lon are 1D, then uncomment this line
# and see "map" call below.
#
# lon2d,lat2d = np.meshgrid(lon,lat)

# Get all variable names on file
var_names = a.variables.keys()

# Loop through each variable and only plot the 2D vars.
for vname in var_names:
  if (len(a.variables[vname].shape) == 2):
    print "plotting " + vname
    var = a.variables[vname][:]
#    print var.min(),var.max()

# create figure, add axes
    fig = plt.figure(figsize=(8,10))
    ax  = fig.add_axes([0.1,0.1,0.8,0.8])
    ax.set_title(vname)

    map = Basemap(projection='cyl',llcrnrlat=lat.min(),urcrnrlat=lat.max(),
                  resolution='l',  llcrnrlon=lon.min(),urcrnrlon=lon.max())
#
# If your lat/lon are 1D, then use lat2d/lon2d here.
#   x, y = map(lon2d,lat2d)
    x, y = map(lon,lat)


    map.drawcoastlines()
    cf = map.contourf(x,y,var)
    cb = map.colorbar(cf,"bottom", size="7%", pad="10%")

#   plt.show()
    plt.savefig("%s/grib2_plot_%s_mpl.png" % (plots_dir,vname))
    plt.close(fig)

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