#----------------------------------------------------------------------
# Read and plots all 2D variables on a GRIB1 file using PyNIO, 
# matplotlib, basemap.
#
# Curvilinear grid, 113 x 151
#----------------------------------------------------------------------
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_grib1'
plots_dir = 'plots_grib1_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 = "ruc.grb"
a   = Nio.open_file("%s/%s" %(dir,filename))
lat = a.variables["gridlat_236"][:]
lon = a.variables["gridlon_236"][:]

# 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())
    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/grib1_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)