#!/bin/usr/python3
import sys
import argparse
import pandas as pd
import numpy as np

###############################################################################
# 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('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
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
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("Not enough data")
    sys.exit

# Sort OBS and HIS
obs = obs[['Flow']].sort_values(by='Flow').reset_index(drop=True)
his = his[['Flow']].sort_values(by='Flow').reset_index(drop=True)
obs = obs.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)

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['Limit'] = pct_rank_qcut(obs['Flow'], n)
obs = obs.groupby('Limit').mean().reset_index()
obs = obs.set_index('Limit').reindex(np.arange(1, n+1)).fillna(method='ffill').reset_index()
obs = obs.fillna(method='bfill').reset_index()

his['Limit'] = pd.qcut(his['Flow'], pcta)
his = his.groupby('Limit').mean().reset_index()
his['Rank'] = pcta[1:] * 100

# Calculate ratio
rdf = his.drop(['Flow'], axis=1)
#ratio.drop(columns=['Flow'], inplace=True)
rdf['Ratio'] = obs['Flow']/his['Flow']
rdf['Floor'] = rdf['Limit'].apply(lambda x: x.split(',')[0][1:]).astype(float)
rdf.at[0, 'Floor'] = 0

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