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


parser = argparse.ArgumentParser()
parser.add_argument('obsf', help='Observational data for reference')
parser.add_argument('simf', help='Historical simulation for training')
parser.add_argument('outf', help='Training outfile')
parser.add_argument('--his_yrs', nargs=2, type=int,
                    help='Historical years', default=[1970,1999])
parser.add_argument('--fut_yrs', nargs=2, type=int,
                    help='Future years', default=[2000,2099])

args = parser.parse_args()


# Read datafile
obsf = args.obsf
simf = args.simf
outf = args.outf
his_yrs = args.his_yrs
fut_yrs = args.fut_yrs

obsf = '/home/disk/rocinante/DATA/temp/WRF/test/kingcounty/9by10_pnnl.nc'
simf = '/home/disk/rocinante/DATA/temp/WRF/test/kingcounty/9by10_ccsm4.nc'

obs = xr.open_dataset(obsf)
sim = xr.open_dataset(simf)
his = sim.sel(year=slice(his_yrs[0], his_yrs[1]))
fut = sim.sel(year=slice(fut_yrs[0], fut_yrs[1]))



# Testing
lobs = obs.sel(x=0, y=0).PREC.values
lhis = his.sel(x=0, y=0).PREC.values
lfut = fut.sel(x=0, y=0).PREC.values

bc_lib.apply_bc(lobs, lhis, lfut, 'prec', 0.09906)



sys.exit()


# Loop through locations

#calc_bc()


sys.exit()
                

# ##############################################################################
# Creates a training table based on ratio of historical and observation data
# binned by 1% interval percentiles
# ##############################################################################

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)