#!/usr/bin/env python
#
# References:
# 1. Matplotlib: https://matplotlib.org/gallery/pie_and_polar_charts/polar_bar.html?highlight=polar%20charts
# 2. PseudoNeCDF: https://pypi.org/project/PseudoNetCDF/
# 3. Windrose:https://windrose.readthedocs.io/en/latest/index.html
# 4. Pandas: https://pandas.pydata.org/
# 5. Numpy: https://numpy.org/
# 6. argparse: https://docs.python.org/3/library/argparse.html
#
# Credits: PseudoNetCDF
#
import numpy as np
from matplotlib import pyplot as plt
import argparse
import pandas as pd

# getting data and arguments
parser = argparse.ArgumentParser("Providing arguments for pollrose plots")
parser.add_argument('--file', dest='file', required=True,
                     default='file', help='Input Excel file')
parser.add_argument('--sname', dest='sname', required=True,
                     default='sname', help='Site name')
parser.add_argument('--bdate', dest='bdate', required=True,
                     default='bdate', help='Begin date 1-1-2018 format')
parser.add_argument('--edate', dest='edate', required=True,
                     default='edate', help='End date 12-31-2018 format')
parser.add_argument('--binwidth', dest='binwidth', type=float,
                     default=45., help='Direction bin widths in degrees.')
parser.add_argument('--max-pct', dest='maxpct', type=float,
                     default=50., help='Maximum percent on plot')
parser.add_argument('--bounds', dest='bounds', type=lambda x: np.array(eval(x)),
                     default=np.array([10, 20, 30, 50, 100, 120, 140, np.inf]), help='Boundaries for pollutant')
parser.add_argument('--outpath', dest='outpath', type=str, required=True,
                     default=50., help='Destination to plot files')
args = parser.parse_args()
parser.epilog = """
Example:
    $ windrose_mpl_infile.py --file /path/filename --bdate bdate --edate edate --max-pct=50 --bounds="[0.5,1,2,3.5,4.5,6.5,np.inf]" --binwidth 22.5 --outpath outdir
Description:
    Create windrose for site between the dates badte and edate. 
"""
ifile1 = args.file
data = pd.read_csv(ifile1, converters = {'start_hour': int}, index_col=False, engine='python')
data['date'] = pd.to_datetime(data['date'])

ws   = list(data['ws'])
wd   = list(data['wd'])

# windrose plot using matplotlib and numpy
def windrose(wd, ws, args):
    plt.rcParams['text.usetex'] = False
    width = args.binwidth
    halfwidth = width / 2.
    fig = plt.figure(figsize=(6,6))
    ax = fig.add_axes([0.1,0.07,0.8,0.8], polar=True)
    ax.set_theta_offset(np.radians(90))
    ax.set_theta_direction(-1)
    ax.set_xticklabels(['N','NE','E','SE','S','SW','W','NW'],size=14)
    ubs = args.bounds[1:]
    lbs = args.bounds[:-1]
    ubcs = plt.cm.jet(np.arange(len(ubs), dtype='f') / ubs.size)
    color_dict = dict(zip(ubs, ubcs))
    mask = np.ma.getmaskarray(ws) | np.ma.getmaskarray(
        wd) | (ws < args.bounds[0]) | (ws > args.bounds[-1])
    ws = np.ma.masked_where(mask, ws).compressed()
    wd = np.ma.masked_where(mask, wd).compressed()
    dl=pd.to_datetime(args.bdate)
    dlabel=dl.year
    fig.suptitle('%s' % (dlabel), size=20)
    thetas = ((np.radians((wd) // width * width)).astype('float64') %
             (2 * np.pi)).astype('float64')
    uthetas = np.unique(thetas)
    nws = float(ws.size)
    for i,theta in enumerate(uthetas):
        tws = ws[thetas == theta]
        for ubi, ub in enumerate(ubs[::-1]):
            ubc = color_dict[ub]
            pct = (tws < ub).sum() / nws * 100.
            ax.bar(theta, pct, width=np.radians(width), bottom=0.0, color=ubc)
    # creating labels
    labels=[]
    for i,l in enumerate(ubs):
        if (i == 0):
            labels.append('< %s' % l)
        elif (i == 'inf'):
            labels.append('> %s' % ubs[i-1])
        else:
            labels.append('%s - %s' % (ubs[i-1],l))
    ## to add legends in the plot, uncomment below lines
    #rectangles = [plt.Rectangle((0, 0), 1, 1, color=ubc) for ubc in ubcs]
    #fig.legend(rectangles, labels, ncol=3,
    #    bbox_to_anchor=(.5, 0.025), loc='lower center')
    #ax.text(0.27, -0.09, 'Wind speed (m/s)', horizontalalignment='center',
    #    verticalalignment='center', transform=ax.transAxes,size=12)
    ## to add date labels in the plot, uncomment below lines
    #datelabel = '%s_%s' % (args.bdate,args.edate)
    #ax.text(0.95, 1.05, datelabel, horizontalalignment='center', style='italic',
    #    verticalalignment='center', transform=ax.transAxes,size=10)
    ticks = np.linspace(0, args.maxpct, 6)[1:]
    labels = ['%s%%' % i for i in ticks]
    plt.yticks(ticks,labels,size=16)
    ax.set_rmax(args.maxpct)
    ax.set_clip_on(False)
    figpath =  '%s/WRose_%s_%s_%s.png' % (args.outpath,args.sname,args.bdate,args.edate)
    fig.savefig(figpath)
    print('Saved fig', figpath)
    plt.close(fig)

windrose(wd[:],ws[:],args)