#!/usr/bin/env python
#

'''
Sample code for plotting time series of ACCESS-S 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 argparse
import matplotlib.pyplot as plt
import numpy as np
import logging
import datetime


def main():
    '''
    This is the basic program to plot 5%, 50% and 95% values for all ensembles
    for a given location and start date.

    Usage:
        plot_ens_all.py var mod_start_date lat_reqd lon_reqd
    e.g.
        plot_ens_all.py tasmax 19900101 -37.4 145

   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
       lat_reqd - latitude required (float)
       lon_reqd - longitude reuired (float)

   Outputs:
        png file with line plots showing the 5%, 50%, 95%

    To do:
        It expects 11 ensembles, each of 217 days length,
            in a future version this will be an optional argument
        add option of user defined logging file
    '''

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

    parser = argparse.ArgumentParser(description='plotting ACCESS-S ensemble data')
    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('lat_reqd', type=float,
                        help='latitude required ')
    parser.add_argument('lon_reqd', type=float,
                        help='longitude required ')
    args = parser.parse_args()

    var = args.var
    mod_start_date = args.mod_start_date
    lat_reqd = args.lat_reqd
    lon_reqd = args.lon_reqd
    n_ens = 11
    ndays = 217

    ensembles = ['e{:02d}'.format(a1) for a1 in np.arange(1, n_ens+1)]

    ens_all = np.zeros((ndays, n_ens))

    for iens, ens in enumerate(ensembles):
        logging.info('in loop, about to call read_single_point')
        ens_all[:, iens], time_real = read_single_point(
            var, mod_start_date, lat_reqd, lon_reqd, ens)

    ttle = '{var} from {mod_start_date}'.format(var=var, mod_start_date=mod_start_date)

    plt.figure()
    bot, = plt.plot(time_real, np.percentile(ens_all[:, :], 5, axis=1), 'b')
    median, = plt.plot(time_real, np.percentile(ens_all[:, :], 50, axis=1), 'g')
    top, = plt.plot(time_real, np.percentile(ens_all[:, :], 95, axis=1), 'r')
    plt.legend((bot, median, top), ('5%', '50%', '95%'))
    plt.title(ttle)
    plt.gcf().show()
    plt.savefig('ens_demo.png')


def read_single_point(var, mod_start_date, lat_reqd, lon_reqd, ens):

    '''
    This will extract the time series for a given point for a single ensemble

    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
       lat_reqd - latitude required (float)
       lon_reqd - longitude reuired (float)
       ens - ensemble number to be read in (string eg e01)

    Returns:
        time series of var at (lat_requd,lon_reqd)
        array of dates corresponding to time series
    '''

    logging.basicConfig(filename='logging.txt', level=logging.WARNING)
    logging.info('working on {} start date {} lat={}, lon={}, ens number {}'.format(
        var, mod_start_date, lat_reqd, lon_reqd, ens))

    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'][:]
    lons = dataset.variables['lon'][:]

    # find closest point
    lat_indx = np.argmin(np.abs(lats - lat_reqd))
    lon_indx = np.argmin(np.abs(lons - lon_reqd))

    # extract a time series:
    tseries = dataset.variables[var][:, lat_indx, lon_indx]

    time_units = dataset.variables['time'].units
    time_num = dataset.variables['time'][:]
    time_real = netCDF4.num2date(time_num, time_units)

    return tseries, time_real

if __name__ == '__main__':
    main()