#!/usr/bin/python3
import xarray as xr
import pandas as pd
import sys
import os
import datetime as dt

## 

args = sys.argv
ifile = args[1]

gcm = os.path.basename(os.path.dirname(ifile))
#odir = "/home/disk/rocinante/DATA/temp/kcp3/forcings/5.forc_ascii/"
#odir = "/home/disk/rocinante/DATA/temp/kcp3/forcings/xtra/5.forc_ascii/"
#odir = "/home/disk/rocinante/DATA/temp/kcp3/forcings/xtra/murphy/forc/"
odir = "/home/disk/becassine/jswon11/5.forc_ascii/"
#pdir = "/home/disk/picea/mauger/2020_12_SnohoCounty_Flooding/DATA/bc_forcs/lapse_rate_4.5/pnnl_wrf/unifm_bc"

gdir = "{}/{}/".format(odir, gcm)
dtfmt = "%m/%d/%Y-%H:%M:%S"
os.makedirs(gdir, exist_ok=True)


fname = os.path.basename(ifile).replace('.nc', '')
#pfile = "{}/{}".format(pdir, fname)
vcols = [ 'T2', 'WND', 'RH', 'SWDOWN', 'GLW', 'PREC' ]
tz = -8
dz = dt.timedelta(hours=tz)

# Read in forcing data
ds = xr.open_dataset(ifile)
ds['PREC'] = ds.PREC / 1000

# Read in pnnl file used for bias-correction
#pf = pd.read_csv(pfile, header=None, sep=' ')
#pf.columns = ['date'] + vcols

# Shift pnnl time to match data
#pf.index = pd.to_datetime(pf.date, format=dtfmt)
#pf.index = pf.index - dz

# Remove leap day and select overlapping data points to calculate average for bc
#pf = pf[~((pf.index.month == 2) & (pf.index.day == 29))][vcols]
#pmean = pf.mean()
#dmean = ds.sel(times = pf.index).mean()

# Apply bias-correction on  temp, wind, sw, and prec
#ds['T2'] = ds.T2 + (pmean.T2 - dmean.T2)
#ds['WND'] = ds.WND * (pmean.WND / dmean.WND)
#ds['SWDOWN'] = ds.SWDOWN * (pmean.SWDOWN / dmean.SWDOWN)
#ds['PREC'] = ds.PREC * (pmean.PREC / dmean.PREC)

# Shift time back to local time
df = ds.to_dataframe()
df = df[['T2', 'WND', 'RH', 'SWDOWN', 'GLW', 'PREC']]
df = df[~((df.index.month == 2) & (df.index.day == 29))]
dates = pd.date_range(df.index[0] + dz, df.index[-1] + dz, freq='H')
dates = dates[~((dates.month == 2) & (dates.day==29))]
df.index = dates

# Save forcing
df['PREC'] = df.PREC.apply(lambda x: "{:0.8f}".format(x))
out = "{}/{}".format(gdir, fname)
df.to_csv(out, header=False, date_format=dtfmt, sep=' ', float_format='%.3f')