#!/usr/bin/env python
#

'''
Sample code for plotting 2d data from OPeNDAP

Date: 27/10/2017

Dependencies: Python 2.7 with numpy, netCDF4, matplotlib

Tested using: python/2.7.6, pythonlib/netCDF4/1.0.4, numpy-1.8.0-py2.7, matplotlib-1.3.1-py2.7-linux-x86_64
'''

__Author__ = "Morwenna Griffiths"


import netCDF4
import matplotlib.pyplot as plt
import numpy as np
import datetime
import logging
import argparse
import shapefile as shp
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from matplotlib.patches import PathPatch

DEFAULT_SHAPEFILE = None


def main():
    '''
    This is the basic program to make a contour plot a 2d variable for a given start date and
    valid time.
    It will also plot a shapefile over the top.

    Usage:
        plot_single_ens.py var mod_start_date ens_num valid_date
    e.g.
        plot_single_ens.py 'tasmax' '19900501' 1 '19900512'
    or
        plot_single_ens.py 'tasmax' '19900501' 1 '19900512' --map_bounds 148. 152. -35 -32.5

   Args:
       var - variable to plot, can be:
           tasmax, tasmin, vprp_09, vprp_15, wind_speed, evap, pr, rsds
       mod_start_date - model start date in format YYYYMMDD
       ens - ensemble number to plot
       time_reqd - valid time for the plot YYYYMMDD

    Optional Args:
       shapefile - shapefile to show on the plot (also defines the lat/lon limit for map
       map_bounds - lat_min, lat_max, lon_min, lon_max

    To Do:
        choose sensible max/min for the sub plot
        check that the bounds are within the data
        decide what to do if both shapefile and map_bounds are set
        add option of user defined logging file
    '''

    logging.basicConfig(filename='logging.txt', level=logging.DEBUG)

    parser = argparse.ArgumentParser(description='fixing compliance files')
    parser.add_argument('var', type=str,
                        help='variable required')
    parser.add_argument('mod_start_date', type=str,
                        help='model start date YYYYMMDD required')
    parser.add_argument('ens', type=int,
                        help='ensemble number required ')
    parser.add_argument('time_reqd', type=str,
                        help='Time required YYYYMMDD ')
    parser.add_argument('--shapefile', type=str, default=DEFAULT_SHAPEFILE,
                        help='optional shapefile')
    parser.add_argument('--map_bounds', type=float, nargs='+',
                        help='map bounds (lat_min, lat_max, lon_min, lon_max)')
    args = parser.parse_args()

    var = args.var
    mod_start_date = args.mod_start_date
    ens = 'e{:02d}'.format(args.ens)
    yyyy = int(args.time_reqd[0:4])
    mm = int(args.time_reqd[4:6])
    dd = int(args.time_reqd[6:8])
    time_reqd = datetime.datetime(yyyy, mm, dd)
    myshapefile = args.shapefile
    map_bounds = args.map_bounds

    data_2d, lats, lons = read_single_time(var, mod_start_date, time_reqd, ens)

    ttle = '{var} at {ftime}'.format(var=var, ftime=time_reqd.strftime("%Y-%m-%d"))

    plt.figure()
    plt.pcolormesh(lons, lats, data_2d)
    plt.colorbar()
    plt.title(ttle)
    plt.savefig('sample_map.png')

    if (myshapefile is not None):
        sf = shp.Reader(myshapefile)
        recs = sf.records()
        shapes = sf.shapes()
        Nshp = len(shapes)
        cns = []
        for nshp in range(Nshp):
            cns.append(recs[nshp][1])
        cns = np.array(cns)

        pts = np.array(shapes[nshp].points)
        prt = shapes[nshp].parts
        par = list(prt) + [pts.shape[0]]
        ptchs = []
        for pij in range(len(prt)):
            ptchs.append(Polygon(pts[par[pij]:par[pij+1]], alpha=1))

    extent = (lons.min(), lons.max(), lats.min(), lats.max())

    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.imshow(data_2d, interpolation='none', extent=extent, vmin=15, vmax=25)
    if (myshapefile is not None):
        ax.add_collection(PatchCollection(ptchs, facecolor='none', edgecolor='k', linewidths=1))
        ax.set_xlim(pts[:, 0].min() - 0.15, pts[:, 0].max() + 0.15)
        ax.set_ylim(pts[:, 1].min() - 0.15, pts[:, 1].max() + 0.15)
    elif (map_bounds is not None):
        logging.info(map_bounds)
        ax.set_xlim(map_bounds[0], map_bounds[1])
        ax.set_ylim(map_bounds[2], map_bounds[3])
    else:
        logging.info('no bounds or shapefile set.  Plotting whole domain')

    plt.colorbar()
    plt.title(ttle)
    plt.savefig('sample_map_zoom.png')


def read_single_time(var, mod_start_date, time_reqd, ens):

    '''
    This will extract the data for a single ensemble at a given time

    Args:
       var - variable to plot, can be:
           tasmax, tasmin, vprp_09, vprp_15, wind_speed, evap, pr, rsds
       mod_start_date - model start date in format YYYYMMDD
       time_reqd - valid time for the plot YYYYMMDD
       ens - ensemble number to be read in (string eg e01)

    Returns:
        2d array of the data
        1d array of the associated latitudes
        1d array of the associated longitudes
    '''

    opendap_dir = 'http://dapds00.nci.org.au/thredds/dodsC/ub7/access-s1/hc/calibrated_5km/atmos'

    fname = '{opendap_dir}/{var}/daily/{ens_num}/da5_{var}_{mod_start_date}_{ens_num}.nc'.format(
        opendap_dir=opendap_dir, var=var, ens_num=ens, mod_start_date=mod_start_date)

    logging.info(fname)
    dataset = netCDF4.Dataset(fname)

    # extract the points:
    lats = dataset.variables['lat'][::-1]
    lons = dataset.variables['lon'][:]

    time_units = dataset.variables['time'].units
    time_num = dataset.variables['time'][:]
    time_real = netCDF4.num2date(time_num, time_units)
    time_indx = np.argmin(np.abs(time_reqd - time_real))

    # extract the data, with the latitude data flipped for easier plotting later
    data_2d = dataset.variables[var][time_indx, ::-1, :]

    return data_2d, lats, lons

if __name__ == '__main__':
    main()