###############################################################################
#
# write_data_uncoupled.R
#
# You are welcome to use this code for your own research but please 
# properly cite and attibute:
#
# Author: 
# Cameron Bracken, cameron.bracken@gmail.com
#
# Citation: 
# C. Bracken, K. Holman, B. Rajagopalan, H. Moradkhani, A Bayesian 
# hierarchical approach to multivariate nonstationary hydrologic frequency 
# analysis, Water Resources Research, 2017, Under Review. 
#
# This file prepares input data for the Bayesian sampler.
# The main inputs are three annual maximum  timeseries (snow, elevation,
# streamflow) for the Fontenelle Dam in Colorado and climate covariate. 
# The covariates are seasonal average (MAM).
#
# The variable names should either be fairly obvious (covars_snow, for the snow covariates) 
# or y=snow, z=flow, r=pool elevation, following the manuscript. 
#
###############################################################################

source('lib.R')

load_packages(c('rstan','ggplot2','data.table','reshape2','ismev'))

model_years = 1981:2015
model_name = 'bayes_joint'

sample_file = sprintf('output/%s/samples/samples',model_name)
dir_create(dirname(sample_file))
dir_create('model_data')

snow_dt = fread('data/station_annmax_SWE_1979-2015.csv')
res = fread('data/annmax_WSE_Inflow_TPD_1963-2015.csv')

snow_mat = data.matrix(dcast.data.table(snow_dt[WY %in% model_years],WY~id,value.var='mWESD')[,WY:=NULL])/1000 #km
snow_complete = snow_mat[,apply(snow_mat,2,function(x)!any(is.na(x)))]

flow = res[WY %in% model_years]$INmax/1000 # kcfs
elev = res[WY %in% model_years]$WSEmax/1000 #km
snow = apply(snow_mat,1,mean,na.rm=TRUE)#snow_mat[,'USS0006K08S']

enso_winter = fread('data/mei_djf.csv')
pdo_winter = fread('data/pdo_djf.csv')
amo_winter = fread('data/amo_djf.csv')

indicies_winter = merge(merge(enso_winter,pdo_winter,by=c('year')),
    amo_winter,by=c('year'))[year %in% model_years]
#indicies = merge(enso,pdo,by=c('year'))[year %in% model_years]
covars_snow = cbind(1,scale(as.matrix(indicies_winter)))

enso_spring = fread('data/mei_mam.csv')
pdo_spring = fread('data/pdo_mam.csv')
amo_spring = fread('data/amo_mam.csv')

indicies_spring = merge(merge(enso_spring,pdo_spring,by=c('year')),
    amo_spring,by=c('year'))[year %in% model_years]
#indicies = merge(enso,pdo,by=c('year'))[year %in% model_years]
covars_flow = cbind(1,scale(as.matrix(indicies_spring)))

#covars_snow = covars_snow[,c(1,4+1)]
#covars_flow = covars_flow[,c(1,3+1)]
#snow4_flow7

np_flow = ncol(covars_flow)
np_snow = ncol(covars_snow)
nt = length(model_years)
ns = ncol(snow_mat)

par_names = c('beta_loc_snow', 'beta_loc_flow', 
    'scale_snow', 'scale_flow', 'scale_elev',
    'shape_snow', 'shape_flow', 'shape_elev',
    'cor_snow_flow', 'cor_snow_elev', 'cor_flow_elev',
    'a','b','c')

par_lens = c(np_snow, np_flow, rep(1,12))
par_dims = list(np_flow,np_snow,1,1,1,1,1,1,1,1,1,1,1,1)
n_pars = sum(par_lens)
names(par_dims) = par_names

init = numeric(n_pars)
init[] = c(c(10,rep(0,np_snow-1)), # betas
        c(10,rep(0,np_flow-1)),
        c(5,5,5), # scale
        c(0.001,0.001,0.001), # shape
        rep(.8,3), # correlations
        c(max(elev),0.002,15) # elevation coefficients
    )

lower=c(rep(-10,np_snow+np_flow), # betas
        c(0,0,0), # scale
        c(-1,-1,-1.5), # shape
        rep(-1,3), # correlations
        c(0,0,1) # elevation coefficients
        )

upper=c(c(10,rep(10,np_snow-1)), # betas
        c(10,rep(10,np_flow-1)),
        c(10,10,10), # scale
        rep(1,3), # shape
        rep(1,3), # correlations
        c(10,0.01,100) # elevation coefficients
        )

save(np_snow,np_flow,nt,par_dims,n_pars,model_years,
    lower,upper,init,snow_mat,snow,flow,elev,
    covars_snow,covars_flow,
    file=sprintf('model_data/%s_data.RData',model_name))

write(c(np_snow, np_flow, nt, n_pars),'model_data/sizes')
write(init, 'model_data/init')
write(snow,'model_data/y')
write(flow,'model_data/z')
write(elev,'model_data/r')
write(covars_snow,'model_data/covars_snow')
write(covars_flow,'model_data/covars_flow')
write(lower,'model_data/lower')
write(upper,'model_data/upper')