#############################################################################
#
# postprocess.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 script is responsible for postprocessing the sampler output. Modify 
# the parameters below to match the ones used when you ran the sampler.  
#
#############################################################################

source('lib.R')
load_packages(c('coda','data.table','reshape2','parallel','compiler',
    'copula','zoo'))

# when thin=n, every n samples will be used. Helps to reduce autocorrelation in the chains
thin = 2
# when warmup=n, will discard the first 
warmup = 1000
# total number of samaples in the sample file
n_samples = 5000
# number of cores to use when processing the output in parallel
cores = 8

# name for the output directory. Should match the one used in write_data.R
model_name = 'bayes_joint'
output_dir = sprintf('output/%s',model_name)
method = 'bayes_joint'

load(sprintf('model_data/%s_data.RData',model_name))

posterior_plot_dir = file.path(output_dir,'plots/posterior_plots')
trace_plot_dir = file.path(output_dir,'plots/trace_plots')
model_plot_dir = file.path(output_dir,'plots/model_plots')
dir_create(c(posterior_plot_dir,trace_plot_dir,model_plot_dir))

par_labs = make_par_names(par_dims)
sample_ind = seq(warmup+1,n_samples,by=thin)

obs = data.table(model_years,flow,elev,snow)

x1 = matrix(scan(file.path(output_dir,'samples/samples_01')),nrow=n_samples,byrow=T)
x2 = matrix(scan(file.path(output_dir,'samples/samples_02')),nrow=n_samples,byrow=T)
x3 = matrix(scan(file.path(output_dir,'samples/samples_03')),nrow=n_samples,byrow=T)

colnames(x1) = par_labs
colnames(x2) = par_labs
colnames(x3) = par_labs

m = mcmc.list(mcmc(x1[sample_ind,]), mcmc(x2[sample_ind,]), mcmc(x3[sample_ind,]))
saveRDS(as.matrix(m),file.path(output_dir,'pars.rds'))

rhat = gelman.diag(m,autoburnin=F)
max_rhat = max(rhat$psrf[,1],na.rm=TRUE)
message(sprintf('The max Rhat is: %4.2f', max_rhat))

posterior_plots(m,posterior_plot_dir)
trace_plots(m,trace_plot_dir)
#summary = summary(m)
############################


############################
lp1 = scan(file.path(output_dir,'samples/lp_01'))#[sample_ind]
lp2 = scan(file.path(output_dir,'samples/lp_02'))#[sample_ind]
lp3 = scan(file.path(output_dir,'samples/lp_03'))#[sample_ind]
lp = data.frame(lp1[sample_ind],lp2[sample_ind],lp3[sample_ind])
lp$iteration = sample_ind
setnames(lp,c('1','2','3','iteration'))
lpdt = data.table(melt(lp,id.vars='iteration',variable.name='chain'))

message(sprintf('The DIC is: %4.2f', dic(c(lp1[sample_ind],lp2[sample_ind],lp3[sample_ind]))))

p_lp = ggplot(lpdt)+geom_line(aes(iteration,value,color=chain))+theme_minimal()
ggsave(sprintf('%s/lp.pdf',trace_plot_dir),p_lp,width=7,height=7)
######################################

cl = makeCluster(cores)

pars_list = apply(as.matrix(m),1,extract_par_groups,par_dims)

return_level_snow <- function(pars, covars, p=0.99){

    library(evd)
    nt = nrow(covars)
    loc_snow = covars %*% pars$beta_loc_snow
    qgev(p, loc_snow, pars$scale_snow, pars$shape_snow)

}

return_level_flow <- function(pars, covars, p=0.99){

    library(evd)
    nt = nrow(covars)
    loc_flow = covars %*% pars$beta_loc_flow
    qgev(p, loc_flow, pars$scale_flow, pars$shape_flow)
}

return_level_elev <- function(pars, covars, p=0.99){

    library(evd)
    nt = nrow(covars)
    #loc_elev = with(pars,a - exp(-covars %*% beta_loc_flow * b))
    loc_elev = with(pars,a/(1+b*c^(-covars %*% beta_loc_flow))) 
    qgev(p, loc_elev, pars$scale_elev, pars$shape_elev)
}

quantile_curve_snow <- function(pars, covars, p=seq(0.001,0.999,by=0.001)){

    library(evd)
    nt = nrow(covars)
    loc_snow = mean(covars %*% pars$beta_loc_snow)
    qgev(p, loc_snow, pars$scale_snow, pars$shape_snow)

}

quantile_curve_flow <- function(pars, covars, p=seq(0.001,0.999,by=0.001)){

    library(evd)
    nt = nrow(covars)
    loc_flow = mean(covars %*% pars$beta_loc_flow)
    qgev(p, loc_flow, pars$scale_flow, pars$shape_flow)

}

quantile_curve_elev <- function(pars, covars, p=seq(0.001,0.999,by=0.001)){

    library(evd)
    nt = nrow(covars)
    loc_elev = mean(with(pars,a/(1+b*c^(-covars %*% beta_loc_flow))))
    qgev(p, loc_elev, pars$scale_elev, pars$shape_elev)

}

copula_sim <- function(pars,covars_flow,covars_snow){
    library(copula)
    nt = nrow(covars_flow)
    cors = with(pars,c(cor_snow_flow,cor_snow_elev,cor_flow_elev))
    cop = normalCopula(param=cors,dim=3,dispstr='un')
    cop_sim = rCopula(nt,cop)

    loc_snow = covars_flow %*% pars$beta_loc_snow
    loc_flow = covars_snow %*% pars$beta_loc_flow
    #loc_elev = with(pars,a - exp(-covars_flow %*% beta_loc_flow * b))
    loc_elev = with(pars,a/(1+b*c^(-covars %*% beta_loc_flow))) 

    snow = mapply(qgev,cop_sim[,1],loc_snow,pars$scale_snow, pars$shape_snow)
    flow = mapply(qgev,cop_sim[,2],loc_flow,pars$scale_flow, pars$shape_flow)
    elev = mapply(qgev,cop_sim[,3],loc_elev,pars$scale_elev, pars$shape_elev)

    cbind(snow,flow,elev)
}

ind_sim <- function(pars,covars_snow,covars_flow){
    library(copula)
    nt = nrow(covars_flow)

    loc_snow = covars_snow %*% pars$beta_loc_snow
    loc_flow = covars_flow %*% pars$beta_loc_flow
    #loc_elev = with(pars,a - exp(-covars_flow %*% beta_loc_flow * b))
    loc_elev = with(pars,a/(1+b*c^(-covars %*% beta_loc_flow))) 

    p_snow = runif(nt)
    p_flow = runif(nt)
    p_elev = runif(nt)

    snow = mapply(qgev,p_snow,loc_snow,pars$scale_snow, pars$shape_snow)
    flow = mapply(qgev,p_flow,loc_flow,pars$scale_flow, pars$shape_flow)
    elev = mapply(qgev,p_elev,loc_elev,pars$scale_elev, pars$shape_elev)

    cbind(snow,flow,elev)
}


Tr = 1/(1-seq(0.001,0.999,by=0.001))

quantile_snow = parSapply(cl, pars_list, quantile_curve_snow, covars_snow)
rownames(quantile_snow)=Tr
colnames(quantile_snow)=1:ncol(quantile_snow)
quantile_dt_snow = as.data.table(melt(quantile_snow))
setnames(quantile_dt_snow,c('period','sim','value'))
bounds_snow = summarise(group_by(quantile_dt_snow,period),
    q50=median(value),
    q95=quantile(value,.95),
    q05=quantile(value,.05))

quantile_flow = parSapply(cl, pars_list, quantile_curve_flow, covars_flow)
rownames(quantile_flow)=Tr
colnames(quantile_flow)=1:ncol(quantile_flow)
quantile_dt_flow = as.data.table(melt(quantile_flow))
setnames(quantile_dt_flow,c('period','sim','value'))
bounds_flow = summarise(group_by(quantile_dt_flow,period),
    q50=median(value),
    q95=quantile(value,.95),
    q05=quantile(value,.05))

quantile_elev = parSapply(cl, pars_list, quantile_curve_elev, covars_flow)
rownames(quantile_elev)=Tr
colnames(quantile_elev)=1:ncol(quantile_elev)
quantile_dt_elev = as.data.table(melt(quantile_elev))
setnames(quantile_dt_elev,c('period','sim','value'))
bounds_elev = summarise(group_by(quantile_dt_elev,period),
    q50=median(value),
    q95=quantile(value,.95),
    q05=quantile(value,.05))

p_quantile_snow = ggplot(bounds_snow)+
    #geom_line(aes(period,value,group=sim),color=grey(.8))+
    geom_line(aes(period,q50))+
    geom_line(aes(period,q95),lty=2)+
    geom_line(aes(period,q05),lty=2)+
    theme_minimal()+
    scale_x_log10(breaks=c(1,10,100,1000))+
    xlab('Return Period')+
    ylab('Snow Return Level')+
    theme(panel.grid=element_blank())+
    ggtitle('Snow')#+
    #ylim(c(0,15))

p_quantile_flow = ggplot(bounds_flow)+
    #geom_line(aes(period,value,group=sim),color=grey(.8))+
    geom_line(aes(period,q50))+
    geom_line(aes(period,q95),lty=2)+
    geom_line(aes(period,q05),lty=2)+
    theme_minimal()+
    scale_x_log10(breaks=c(1,10,100,1000))+
    xlab('Return Period')+
    ylab('Flow Return Level')+
    theme(panel.grid=element_blank())+
    ggtitle('Flow')#+
    #ylim(c(0,6))

p_quantile_elev = ggplot(bounds_elev)+
    #geom_line(aes(period,value,group=sim),color=grey(.8))+
    geom_line(aes(period,q50))+
    geom_line(aes(period,q95),lty=2)+
    geom_line(aes(period,q05),lty=2)+
    theme_minimal()+
    scale_x_log10(breaks=c(1,10,100,1000))+
    xlab('Return Period')+
    ylab('Elev Return Level')+
    theme(panel.grid=element_blank())+
    ggtitle('Elevation')#+
    #ylim(c(9.30,9.34))

pdf(file.path(model_plot_dir,'quantile_curves.pdf'),width=7,height=2.5)
multiplot(p_quantile_snow, p_quantile_flow, p_quantile_elev,cols=3)
dev.off()


bounds_snow$variable = 'snow'
bounds_flow$variable = 'flow'
bounds_elev$variable = 'elev'
bounds_snow$method = method
bounds_flow$method = method
bounds_elev$method = method

rl = rbind(bounds_snow,bounds_flow,bounds_elev)
saveRDS(rl,file.path(output_dir,'rl.rds'))


###############################################
# 100 year return period 
###############################################
rl_snow = parSapply(cl, pars_list, return_level_snow, covars_snow)
rownames(rl_snow) = model_years
colnames(rl_snow) = 1:ncol(rl_snow)
rl_dt_snow = data.table(melt(rl_snow))
setnames(rl_dt_snow,c('year','sim','value'))

rl_flow = parSapply(cl, pars_list, return_level_flow, covars_flow)
rownames(rl_flow) = model_years
colnames(rl_flow) = 1:ncol(rl_flow)
rl_dt_flow = data.table(melt(rl_flow))
setnames(rl_dt_flow,c('year','sim','value'))

rl_elev = parSapply(cl, pars_list, return_level_elev, covars_flow)
rownames(rl_elev) = model_years
colnames(rl_elev) = 1:ncol(rl_elev)
rl_dt_elev = data.table(melt(rl_elev))
setnames(rl_dt_elev,c('year','sim','value'))

p_rl_snow = ggplot(rl_dt_snow)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    #ylim(c(2,7))+
    ylab('100 year return level')+
    theme_minimal()+
    ggtitle('Snow')
ggsave(file.path(model_plot_dir,'return_level_snow_2yr.pdf'),p_rl_snow,width=6,height=5,dpi=600)

p_rl_flow = ggplot(rl_dt_flow)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    #ylim(c(0,3))+
    ylab('100 year return level')+
    theme_minimal()+
    ggtitle('Flow')
ggsave(file.path(model_plot_dir,'return_level_flow_2yr.pdf'),p_rl_flow,width=6,height=5,dpi=600)

p_rl_elev = ggplot(rl_dt_elev)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    #ylim(c(9.31,9.335))+
    ylab('100 year return level')+
    theme_minimal()+
    ggtitle('Elevation')
ggsave(file.path(model_plot_dir,'return_level_elev_2yr.pdf'),p_rl_elev,width=6,height=5,dpi=600)


rl_dt_snow$variable = 'snow'
rl_dt_flow$variable = 'flow'
rl_dt_elev$variable = 'elev'
rl_dt_snow$method = method
rl_dt_flow$method = method
rl_dt_elev$method = method

rl_ns = rbind(rl_dt_snow,rl_dt_flow,rl_dt_elev)
saveRDS(rl_ns,file.path(output_dir,'rl_ns.rds'))

###############################################
# 2 year return period 
###############################################
rl_snow_2yr = parSapply(cl, pars_list, return_level_snow, covars_snow, p=0.5)
rownames(rl_snow_2yr) = model_years
colnames(rl_snow_2yr) = 1:ncol(rl_snow_2yr)
rl_dt_snow_2yr = data.table(melt(rl_snow_2yr))
setnames(rl_dt_snow_2yr,c('year','sim','value'))

rl_flow_2yr = parSapply(cl, pars_list, return_level_flow, covars_flow, p=0.5)
rownames(rl_flow_2yr) = model_years
colnames(rl_flow_2yr) = 1:ncol(rl_flow_2yr)
rl_dt_flow_2yr = data.table(melt(rl_flow_2yr))
setnames(rl_dt_flow_2yr,c('year','sim','value'))

rl_elev_2yr = parSapply(cl, pars_list, return_level_elev, covars_flow, p=0.5)
rownames(rl_elev_2yr) = model_years
colnames(rl_elev_2yr) = 1:ncol(rl_elev_2yr)
rl_dt_elev_2yr = data.table(melt(rl_elev_2yr))
setnames(rl_dt_elev_2yr,c('year','sim','value'))

p_rl_snow = ggplot(rl_dt_snow_2yr)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    ylim(c(2,7))+
    geom_point(aes(model_years,snow),color='steelblue',data=obs)+
    ylab('2 year return level')+
    theme_minimal()+
    ggtitle('Snow')
ggsave(file.path(model_plot_dir,'return_level_snow_2yr.pdf'),p_rl_snow,width=6,height=5)

p_rl_flow = ggplot(rl_dt_flow_2yr)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    ylim(c(0,3))+
    geom_point(aes(model_years,flow),color='steelblue',data=obs)+
    ylab('2 year return level')+
    theme_minimal()+
    ggtitle('Flow')
ggsave(file.path(model_plot_dir,'return_level_flow_2yr.pdf'),p_rl_flow,width=6,height=5)

p_rl_elev = ggplot(rl_dt_elev_2yr)+
    geom_boxplot(aes(year,value,group=year),outlier.shape=NA)+
    ylim(c(9.31,9.335))+
    geom_point(aes(model_years,elev),color='steelblue',data=obs)+
    ylab('2 year return level')+
    theme_minimal()+
    ggtitle('Elevation')
ggsave(file.path(model_plot_dir,'return_level_elev_2yr.pdf'),p_rl_elev,width=6,height=5)

grDevices::png(file.path(model_plot_dir,'return_level_2yr.png'),width=6,height=10,res=600,unit='in')
multiplot(p_rl_snow,p_rl_flow,p_rl_elev,cols=1)
dev.off()


###########################################################
###########################################################

# sim = parLapply(cl, pars_list, copula_sim, covars_snow, covars_flow)
# sim_cop_dt = as.data.table(do.call('rbind',sim))
# sim_cop_snow = sapply(sim,'[',1:nt,1)
# sim_cop_flow = sapply(sim,'[',1:nt,2)
# sim_cop_elev = sapply(sim,'[',1:nt,3)

# sim_ind = parLapply(cl, pars_list, ind_sim, covars_snow, covars_flow)
# sim_ind_dt = as.data.table(do.call('rbind',sim_ind))
# sim_ind_snow = sapply(sim_ind,'[',1:nt,1)
# sim_ind_flow = sapply(sim_ind,'[',1:nt,2)
# sim_ind_elev = sapply(sim_ind,'[',1:nt,3)

# p_snow_flow_ind = ggplot(sim_ind_dt)+
#     geom_density_2d(aes(snow,flow))+theme_minimal()+
#     geom_point(aes(snow,flow),data=data.frame(snow,flow))
# p_snow_elev_ind = ggplot(sim_ind_dt)+
#     geom_density_2d(aes(snow,elev))+theme_minimal()+
#     geom_point(aes(snow,elev),data=data.frame(snow,elev))
# p_flow_elev_ind = ggplot(sim_ind_dt)+
#     geom_density_2d(aes(flow,elev))+theme_minimal()+
#     geom_point(aes(flow,elev),data=data.frame(flow,elev))
# p_snow_flow_cop = ggplot(sim_cop_dt)+
#     geom_density_2d(aes(snow,flow))+theme_minimal()+
#     geom_point(aes(snow,flow),data=data.frame(snow,flow))
# p_snow_elev_cop = ggplot(sim_cop_dt)+
#     geom_density_2d(aes(snow,elev))+theme_minimal()+
#     geom_point(aes(snow,elev),data=data.frame(snow,elev))
# p_flow_elev_cop = ggplot(sim_cop_dt)+
#     geom_density_2d(aes(flow,elev))+theme_minimal()+
#     geom_point(aes(flow,elev),data=data.frame(flow,elev))
# layout = rbind(c(1,2,3),c(4,5,6))

# pdf(file.path(model_plot_dir,'joint_contours.pdf'),width=7,height=4)
# multiplot(p_snow_flow_ind,p_snow_elev_ind,p_flow_elev_ind,
#           p_snow_flow_cop,p_snow_elev_cop,p_flow_elev_cop,layout=layout)
# dev.off()