#!/usr/bin/python

"""
Python code to read an H5 observation file

Usage:
python read_h5.py -i [innovar_file_name]

"""

# some standard libraries
import os
import sys
# the HDF5 library
import h5py
import tables
import numpy as np
# the MatLab like plotting
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
# the geographic mapping database
from mpl_toolkits.basemap import Basemap
# debugging utility
from IPython import embed as shell

font_base = 12
cm   = matplotlib.cm.get_cmap('bwr')


def main(inputfile, image_dir=None):


    #  Read data file
    data = read_h5_data(inputfile)

    for ich in np.arange(15)+1:
        ch = ('%.2i' % ich)
        title = 'AMSU-A Observation Channel' + ch
        image_name = '_'.join( ['AMSU-A', 'Observation', 'Ch', ch ] )
        image_name = os.path.join(image_dir, image_name)
        ch_loc = data['Channel'] == float(ich)
        if sum(ch_loc) == 0: continue
        plot_global(data['Observation'][ch_loc], 
                    data['Latitude'][ch_loc], 
                    data['Longitude'][ch_loc], 
                    title, image_name)


def read_h5_data(inputfile):


    h5file = h5py.File(inputfile, 'r')   # 'r' means that hdf5 file is open in read-only mode
    #h5file = tables.open_file(inputfile, mode = "r")
    #shell()
    #sys.exit()
    #print h5file.keys()                 # use with h5py
    #print h5file['INNOV'].dtype.names   # use with h5py
    #print h5file.walk_groups            # use with "tables"
    data = {}
    data['Satellite_Id']       = np.array(h5file['INNOV']['satid'])
    data['Sensor_Id']          = np.array(h5file['INNOV']['sensor'])
    data['Obs_Time']           = np.array(h5file['INNOV']['idt'])
    data['Elevation']          = 800.0  # h5file['INNOV']['']
    data['Latitude']           = np.array(h5file['INNOV']['lat'])
    data['Longitude']          = np.array(h5file['INNOV']['lon'])
    data['Satellite_Zenith']   = np.array(h5file['INNOV']['p_ob'])
    data['Solar_Zenith']       = 0.0    # h5file['INNOV']['None']
    data['Channel']            = np.array(h5file['INNOV']['channel'])
    data['Sensor']             = np.array(h5file['INNOV']['sensor'])
    data['Satellite']          = np.array(h5file['INNOV']['satid'])
    data['Scan_Position']      = np.array(h5file['INNOV']['idp'])
    data['Observation']        = np.array(h5file['INNOV']['ob'])
    data['QC_Flag']            = np.array(h5file['INNOV']['ichk'])
    data['Observation_Error']  = np.array(h5file['INNOV']['err'])
    data['Land_Type_Index']    = np.array(h5file['INNOV']['lsi'])

    h5file.close()

    return data

def plot_global(data, lat, lon, title, image_name, clabel=None, range=None):

    # plot the data
    font_base = 12

    fig = plt.figure(figsize=(10.5,7)) #inches at 100 ppi

    # the gridspec nrows and ncols, then ratios for each row
    gs = gridspec.GridSpec(4,1,height_ratios=[48,8,8,3],hspace=0)

    # define plotting axis
    ax = plt.subplot(gs[0])

    # define the map object
    m =  Basemap(projection='cyl', llcrnrlat=-90, urcrnrlat=90,
                                   llcrnrlon=-180, urcrnrlon=180,
                                   resolution='c')
    m.drawparallels(np.arange(-90, 91, 30), labels=[1,0,0,1])
    m.drawmeridians(np.arange(-180, 181, 60), labels=[1,0,0,1])

    # data must all be same type (numpy arrays)
    lat = np.array(lat, dtype='float64' )
    lon = np.array(lon, dtype='float64' )
    data = np.array(data, dtype='float64' )
    lon[lon > 180] -= 360.

    # optional mask outside given range
    if not range:
        range = [ data.min(), data.max() ]
    else:
        data = np.ma.masked_outside( data, range[0], range[1])
        lat = np.ma.array( lat, mask = data.mask )
        lon = np.ma.array( lon, mask = data.mask )

    # transform lon and lat into mappable coordinates
    #x, y = np.meshgrid(lon, lat)
    x, y = m(lon, lat)
    # nomalize data to colormapping
    norm = matplotlib.colors.Normalize(vmin=range[0], vmax=range[1])
    # scatter plot of data
    #shell()
    #sys.exit()
    m.scatter(x, y, s=5, marker="o", c=data, norm=norm, edgecolors='none', cmap=cm)
    m.drawcoastlines()
    # add a title with some stats from data
    data_stats = ( "min=%10.4g    max=%10.4g    mean=%10.4g  stdv=%10.4g" %
                   (data.min(), data.max(), data.mean(), data.std()) )
    print data_stats
    plt.title(data_stats, fontweight='demibold', fontsize=1.0*font_base)

    #  add a colorbar
    histogram_colorbar(data, gs, range=range, clabel=clabel)

    # add an overarching title to figure
    fig.suptitle(title, fontweight='demibold', fontsize=1.5*font_base)

    # write image to file
    fig.savefig('%s.png' % image_name, format='png')
    plt.close(fig)

def histogram_colorbar(data, gs, range=None, clabel=None, bins=66):

    if data is None:
        print "missing data input"
        sys.exit()
    # Use 1-D data for the histogram
    data = data.flatten()

    if range is None:
        N, bins = np.histogram(data,bins=bins)
        bin_frac = 1000. * np.array(N,dtype='float64')/np.float(np.max(N))
        hrange = np.where(bin_frac > 1)
        range = ( bins[np.min(hrange)], bins[np.max(hrange)+1] )
        #range = ( data.min(), data.max() )
    if range[0] == range[1]:
        print "min and max values are equal"
        sys.exit()

    # add a histogram
    ax2 = plt.subplot(gs[2])
    N, bins, patches = ax2.hist(data,bins,range=range)
    sum = np.sum( N )

    # color the histogram bars according to value
    norm = matplotlib.colors.Normalize(vmin=range[0], vmax=range[1])
    bin_centers = 0.5*(bins[1:] + bins[:-1])
    for (bin_center, patch) in zip(bin_centers, patches):
        patch.set_facecolor(cm(norm(bin_center)))
    ax2.set_xlim(range[0],range[1])
    ax2.autoscale(enable=True,axis='y')
    ax2.xaxis.set_major_locator(matplotlib.ticker.NullLocator())
    ax2.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
    ax2.axis('off')
    ax2.set_title('Total: ' + str(sum) )

    # add traditional colorbar too
    cax = plt.subplot(gs[3])
    c = matplotlib.colorbar.ColorbarBase(cax, orientation='horizontal', norm=norm, cmap=cm)
    c.set_label(clabel)


if __name__ == "__main__":

    from optparse import OptionParser

    usage = 'usage: %prog -i input-loc -o image-dir' 

    parser = OptionParser(usage)
    parser.add_option('-i', '--input-file', dest='input_file',
                      action='store', default=None,
                      help='Input files')
    parser.add_option('-o', '--image-dir', dest='image_dir',
                      action='store', default=os.getcwd(),
                      help='Output location for image')

    (options, args) = parser.parse_args()

    # check for presence of input file argument
    if not options.input_file:
        parser.error("no input file given, specify with -i option")
    # check for existence of input file
    if not os.path.isfile(options.input_file):
        parser.error("Missing input file: %s"% options.input_file)

    # if output directory does not exist create it
    if not os.path.isdir(options.image_dir):
        os.mkdirs(options.image_dir)

    main(options.input_file, image_dir=options.image_dir)
