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

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(time=slice(str(his_yrs[0]), str(his_yrs[1])))
fut = sim.sel(time=slice(str(fut_yrs[0]), str(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)

bchis = his.copy()
bcfut = fut.copy()

x = len(his.lat)
y = len(his.lon)
n = (x+1) * 10 + y + 1
var = 'pr'

for i in range(0, x):
    for j in range(0, y):
        ilat = his.lat[i].data
        jlon = his.lon[j].data
        print("{} - {} : {}/{}".format(ilat, jlon, i*10+j+1, n))

        his_ij = his.sel(lat=ilat, lon=jlon)[var]
        fut_ij = fut.sel(lat=ilat, lon=jlon)[var]
        
        (h, f) = bc.apply_bc(obs.pr, his_ij, fut_ij, 'prec', 0.09906)
        bchis[var].data[i, j] = h
        bcfut[var].data[i, j] = f


print('---')
print(type(h))
print(type(bchis['pr']))
print(h.shape)
print(bchis['pr'].shape)

ds = xr.concat([bchis, bcfut], 'time')
ds.to_netcdf(outf)
print(outf)