#!/bin/usr/python3
import os.path
import sys
import argparse
import pandas as pd
import numpy as np
from datetime import datetime as dt
import calendar as cal

###############################################################################
# Creates a training table based on ratio of historical and observation data
# binned by 1% interval percentiles
###############################################################################
parser = argparse.ArgumentParser()
parser.add_argument('obs', help='Observational data for reference')
parser.add_argument('his', help='Historical simulation for training')
parser.add_argument('--win', action='store', type=int, default=20, 
                    help='Seasonal window size. ')
parser.add_argument('out', help='Training outfile')
args = parser.parse_args()

hdr = ['Year', 'Month', 'Day', 'Flow']
obs_file = args.obs
his_file = args.his
out_file = args.out
wdw = args.win
naList = ["NA", "NaN", ""]

def print_full(x):
    pd.set_option('display.max_rows', len(x))
    print(x)
    pd.reset_option('display.max_rows')    
    
# ----------------------------------------------------------------------------
# Read OBS, HIS
print("---: " + obs_file.split('/')[-1])

if(not os.path.isfile(obs_file)):
    print("OBS file does not exist: " + obs_file)
    sys.exit()
if(not os.path.isfile(his_file)):
    print("HIS file does not exist: " + his_file)
    sys.exit()
    
obs = pd.read_csv(obs_file, sep="[-, /\t]*", engine='python',
                  header=None, na_values=naList)
obs.columns = hdr
his = pd.read_csv(his_file, sep="[-, /\t]*", engine='python',
                  header=None, na_values=naList)
his.columns = hdr

# Filter OBS to 1950~2017
obs = obs[(obs['Year'] >= 1950) & (obs['Year'] <= 2017)]

# Filter HIS to xxxx~2005
his = his[(his['Year'] <= 2005)]

# Data length check
if (len(obs) == 0) | (len(his) == 0):
    print(obs_file)
    print("OBS: " + str(len(obs)))
    print(his_file)
    print("HIS: " + str(len(his)))
    print("Not enough data")
    sys.exit()
    
# ----------------------------------------------------------------------------
# Training Algorigthm

# Calculate day of year value 
def get_dofy(df):
    # Calculate datetime
    dates = pd.to_datetime(df['Year'].astype(str) + '/' +
                           df['Month'].astype(str) + '/' +
                           df['Day'].astype(str),
                           format='%Y/%m/%d')
    # Convert to day of year
    df['dofy'] = dates.apply(lambda x: dt.strftime(x, '%j')).astype(int)

    # Adjust for leap day so all years start on 1 and end on 366. 60 used for leap day
    cor = df['Year'].apply(lambda x: 0 if cal.isleap(x) else 1)
    cor[df['dofy'] < 60] = 0
    df['dofy'] += cor
    return df


# Calculate percentile ranking
def pct_rank_qcut(series, n):
    edges = pd.Series([float(i) / n for i in range(n + 1)])
    f = lambda x: (edges >= x).values.argmax()
    return series.rank(pct=1).apply(f)


obs = get_dofy(obs)
his = get_dofy(his)
rdf = pd.DataFrame()

# Calulate ranking with moving window
for i in range(1,367):

    # Filter only days within range
    obs['var'] = (obs['dofy'] - i + wdw) % 366
    his['var'] = (his['dofy'] - i + wdw) % 366
    obsT = obs[obs['var'] < 2 * wdw + 1].copy()
    hisT = his[his['var'] < 2 * wdw + 1].copy()

    # Sort OBS and HIS
    obsT = obsT[['Flow']].sort_values(by='Flow').reset_index(drop=True)
    hisT = hisT[['Flow']].sort_values(by='Flow').reset_index(drop=True)
    obsT = obsT.dropna()
    
    # Bin OBS and HIS into percentile bins with min streamflow
    n = 100
    pcta = np.arange(0.0, 1.0 + 1.0/n, 1.0/n)
    obsT['Limit'] = pct_rank_qcut(obsT['Flow'], n)

    obsT = obsT.groupby('Limit').mean().reset_index()
    obsT= obsT.set_index('Limit').reindex(np.arange(1,n+1)).fillna(method='ffill').reset_index()
    obsT= obsT.fillna(method='bfill').reset_index()
    #obsT['Rank'] = pcta[1:]

    x = pd.qcut(hisT['Flow'], pcta, retbins=True)[1][:-1]
    hisT['Limit'] = pd.qcut(hisT['Flow'], pcta)
    hisT = hisT.groupby('Limit').mean().reset_index()
    hisT['Rank'] = pcta[1:] * 100

    # Calculate ratio
    rdfT = hisT.drop(['Flow'], axis=1)
    rdfT['Ratio'] = obsT['Flow']/hisT['Flow']
    #rdfT['Floor'] = rdfT['Limit'].apply(lambda x: x.split(',')[0][1:]).astype(float)
    rdfT['Floor'] = x
    rdfT.at[0, 'Floor'] = 0
    rdfT['dofy'] = i
    rdf = rdf.append(rdfT)    
    #print('Process day: ' + str(i))

# Write out training table to file
rdf.to_csv(out_file, index=False)