####################################################################################
# MFA Uncertainty Analysis: Bootstrapping the Sediment Yield Curve
# to create boot-strapped uncertainty intervals for Qeff and Qs50
# Joel Sholtes, Colorado State University
# April, 2014
####################################################################################

## Pre-Amble: Load Functions ###############################################################

rm(list = ls())
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty")
# path <- "bedload_sites.csv"
library(boot) # bootstrap
library(zoo) # rollmean
source("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/minor.ticks.axis.R")
#load Finney (1941) function to carry out Cohn et al. (1989) MVUE Bias Corr. Factor calc
source("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Finney.R")
library(coda) # for effective sample size function
library(MASS) # RLM

# Custom Functions

# Create trapezoidal rule function
trapz <- function (x, y) {
  idx = 2:length(x)
  return (as.double( (x[idx] - x[idx-1]) %*% (y[idx] + y[idx-1])) / 2)
}
# Function: opposite of %in% or match
'%!in%' <- function(x,y)!('%in%'(x,y)) 

# Cumulative Trapezoidal Integration
cum.int <- function (x, y) {
  idx = 2:length(x)
  areas = (x[idx] - x[idx-1]) * (y[idx] + y[idx-1])/ 2
  return (cumsum(areas)/sum(areas))
}

## Import Data ##############################################################
## Sediment Data (Suspended Sites)
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Bedload")

## Import flow file data
path <- "site_summary.csv"
flow.sites <- read.table(path, header = TRUE, sep = ",", 
                         stringsAsFactors = F)
flow.sites <- flow.sites[order(flow.sites$site.no),]


# Read in Bed Load Sed Data
sed.data <- read.csv("BedLoad_R_Data.csv", stringsAsFactors = F, header = TRUE)

# Create empty data frame for results and assign col names to it
cnames <- c("site.no","site.nm", 
            "n.sed","q.min.sed","q.max.sed",
            "b0.low","b0","b0.upp","b1.low","b1","b1.upp",
            "bcf.qmle.ols","bcf.mvue", "r2","sigma.e",
            "rl", "n.eff", "bw", "max.var.kdf", "q.min.flow","q.max.flow","cv", "flash.RB",
            "qh.low", "qh", "qh.upp", "spread.qh",
            "qe.low", "qe", "qe.upp", "spread.qeff")

results <- data.frame(matrix(data = 0, nrow = nrow(flow.sites), ncol = length(cnames)),
                      row.names = NULL, stringsAsFactors = F)
names(results) <- cnames  # Assigns col names to data frame

## Loop Uncertainty Analysis ############################################################

# Loop to conduct regression, rating curve, and CDF / PDF uncertainty analysis 
# and calc / record independent variable metrics (e.g., sigma.e and flow CV and RL) 

for (i in 1:nrow(flow.sites)){
  
  # Extract sed data for regression
  site.data <- subset(sed.data, site.id == flow.sites$site.no[i])
  n <- nrow(site.data)
  
  # Import Flow Data
  path <- paste("DV_Flows/", flow.sites$site.no[i], ".csv", sep = "")
  flow.data <- read.table(path, header = T, stringsAsFactors = F, sep = ",")
  flow <- flow.data$val*0.028316847 # Convert to cms
  
  # Calc Flow Metrics
  N <- length(flow)
  rl <- floor(N/365)
  cv <- round(sqrt(var(flow, na.rm = T)) / mean(flow, na.rm = T),2)
  flash.RB <- round(sum(abs(diff(flow)))/sum(flow) , 3)
  q.min.flow <- signif(min(flow, na.rm= T),2)
  q.max.flow <- round(max(flow, na.rm= T),2)
  
  # Extract Site specific flow - concentration pairs and create temp  
  # x = log(Q[cms]), y = log(Qs[tonnes/day]) vectors for regression
  x <- site.data$cms
  q.min.sed <- signif(min(x, na.rm= T),2)
  q.max.sed <- round(max(x, na.rm= T),2)
  y <- site.data$kg.sec*86.4 # 864 converts from kg/sec to tonnes/day
       
  # Write data to table along with CI for regression line
  sed.site <- data.frame(cms = x, t.day = y, log.q = log10(x), log.qs = log10(y))
  
  # now remove zero values for log-linear regression
  if (min(y) == 0) {
    sed.site <- sed.site[which(y != 0),]
  }
  
  # OLS Log-Linear Regression, needs to call from a data.frame for predict() to work
  model.ols <- lm(log.qs ~ log.q, data = sed.site)
  b0 <- as.numeric(round(coef(model.ols)[1], 3))
  b1 <- as.numeric(round(coef(model.ols)[2], 3))
  bcf.qmle <- signif(exp(sum( residuals(model.ols)^2 )/(n-2) /2), 3)
  r2 <- signif(summary(model.ols)$r.squared, 2)
  sigma.e <- signif(summary(model.ols)$sigma, 2)
  ## Create Confidence & Prediction Bands 
  ## as well as 95% Rating Curves from the CI's of the regression params
  
  # new sequence of flows for rating curve, variable name in data.frame must match
  # that used in the model.ols
  
  # vector of q's for rating curve
  # NOTE: factor of 1.1 is used for min side of vector. 
  # This only works if min(log(flow)) < 1. Might need an if-else statement here'
  min.flow <- min(flow[which(flow > 0)])
  
  x.pred <- if(min.flow < 1) {
    data.frame(log.q = seq(log10(min.flow)*1.1, max(log10(flow))*1.1, 0.005))
  }else{
    data.frame(log.q = seq(log10(min.flow)*.9, max(log10(flow))*1.1, 0.005))
  }
 
  ## Calc upper and lower rating curves based on 95% CI's for params
  
  # CIs for model params
  CI <- confint(model.ols, level = 0.95)
  
  # concatenate lower, estimated, and upper 95% CI for model params (log-transformed)
  b0.ci  <- c(CI[1,1], b0, CI[1,2])
  b1.ci  <- c(CI[2,1], b1, CI[2,2])
  
  ###########################################################
  
  ######################
  ### Log-Linear RLM ###
  ######################
  
  data <- data.frame(log.q = sed.site$log.q, log.qs = sed.site$log.qs)
  R = 1000 # Number of Bootstrap samples
  x.range = range(data$log.q)
  
  # x-vector for linear regression mean response
  x1 <- seq(x.range[1], x.range[2], length.out = 100) 
  
  # original OLS model
  model.ols <- lm(log.qs ~ log.q, data = data)
  bhat = coef(model.ols)
  
  ## Set up Bootstrap Sampling for CI creation: Case Re-sampling for the mean response
  # Davison and Hinkley: 6.5 p: 307-315
  # Create function for bootstrap sampling
  bs <- function(data, indices, x) {
    # Allow boot to select sample 
    data.boot <- data[indices,]
    fit.boot <- lm(data.boot[,2] ~ data.boot[,1])
    return(c(coef(fit.boot))) 
  } 
  
  boot.out <- boot(data, statistic = bs, R = R, parallel = "snow")
  b0.boot <- 10^boot.out$t[,1]
  b1.boot <- boot.out$t[,2]
  
  ###########################################################
  # Calculate parameters for re-transformation
  logQ_bar <- mean(data$log.q)
  m = n - 2 # degrees of freedom
  # variance of discharges in regression dataset
  Sxx = sum((sed.site$log.q - logQ_bar)^2) 
  
  # Add vector of Q in cms
  x.pred <- data.frame(x.pred, cms = 10^(x.pred$log.q))
  
  # Calculate MVUE BCF ->  see Cohn et al. (1989)
  # Variance of new compared with regression dataset Q
  var = 1/n + (x.pred$log.q - logQ_bar)^2 / Sxx
  bcf = numeric(nrow(x.pred))
  
  for (j in 1:nrow(x.pred) ) {
    bcf[j] = Finney(m, (m + 1) / (2 * m) * (1 - var[j]) * sigma.e^2)
  }
  bcf.mvue <- round(max(bcf),2)
  
  #### Kernel Density Function and Normal Approx CB ###
  # Effective Sample Size
  n.eff <- round( as.numeric( effectiveSize(as.mcmc(flow)) ),0 )
  
  kappa <- 1/2/sqrt(pi) # int(K(x)^2 dx) from x = (-inf, inf) for Gaussian kern
  
  bw = mean(c(bw.nrd0(flow),bw.nrd(flow),bw.SJ(flow)))
  
  kde <- density(flow, n = 1028, from = 0, bw = bw)
  
  h <- kde$bw
  
  ############# Normal Approx for Conf Bands ####################
  # Create Naive Confidence Bounds
  alpha <- 0.05
  z.alpha <- -qnorm(p = alpha/2)
  var.kdf <- sqrt(1/(n.eff*h) * kappa^2 * kde$y)
  max.var.kdf <- max(var.kdf)
  kde.cb1 <- rbind(kde$y - var.kdf*z.alpha,
                   kde$y + var.kdf*z.alpha)
  # Replace negative values with zero
  kde.cb1[which(kde.cb1 < 0, arr.ind = T )] <- 0  
  
  # Conf Bands for KDF using x.pred vector in arithmatic space
  kde.cb2 <- rbind(approx(x = kde$x, y = kde.cb1[1,], xout = x.pred$cms)$y, # Lower
                   approx(x = kde$x, y = kde$y, xout = x.pred$cms)$y, # Mean
                   approx(x = kde$x, y = kde.cb1[2,], xout = x.pred$cms)$y # Upper
  )
  kde.cb2[which(is.na(kde.cb2), arr.ind = T)] <- 0 # remove NA values
  

  ###############################################################################
  # Calc booted Sediment Yield Curve
  eff.boot.l <- matrix(ncol = R, nrow = nrow(x.pred))
  eff.boot.m <- matrix(ncol = R, nrow = nrow(x.pred))
  eff.boot.u <- matrix(ncol = R, nrow = nrow(x.pred))
  for(j in 1:R){
    eff.boot.l[,j] <- b0.boot[j]*bcf*x.pred$cms^b1.boot[j] * kde.cb2[1,]
    eff.boot.m[,j] <- b0.boot[j]*bcf*x.pred$cms^b1.boot[j] * kde.cb2[2,]
    eff.boot.u[,j] <- b0.boot[j]*bcf*x.pred$cms^b1.boot[j] * kde.cb2[3,]
  }
  
  eff.l <- apply(eff.boot.l, 1, quantile, probs = 0.025)
  eff.m <- apply(eff.boot.m, 1, quantile, probs = 0.5)
  eff.u <- apply(eff.boot.u, 1, quantile, probs = 0.975)

  Max.eff.boot.l <- numeric(R)
  Max.eff.boot.u <- numeric(R)
  idx.l <- numeric(R)
  idx.m <- numeric(R)
  idx.u <- numeric(R)
  for(j in 1:R){
    idx.l[j] <- which.max(eff.boot.l[,j])
    idx.m[j] <- which.max(eff.boot.m[,j])
    idx.u[j] <- which.max(eff.boot.u[,j])
    Max.eff.boot.l[j] <- eff.boot.l[idx.l[j],j]
    Max.eff.boot.u[j] <- eff.boot.u[idx.l[j],j]
  }
  Qeff.boot.l <- x.pred$cms[idx.l]
  Qeff.boot.m <- x.pred$cms[idx.m]
  Qeff.boot.u <- x.pred$cms[idx.u]
  
  # Calculate cumulative integral of eff for cum sed yield curves
  x.c <- rollmean(x.pred$cms, k = 2)
  cum.boot.l <- apply(eff.boot.l, 2, FUN = cum.int, x = x.pred$cms)
  cum.boot.m <- apply(eff.boot.m, 2, FUN = cum.int, x = x.pred$cms)
  cum.boot.u <- apply(eff.boot.u, 2, FUN = cum.int, x = x.pred$cms)
  
  Qs50.boot.l <- numeric(R)
  Qs50.boot.m <- numeric(R)
  Qs50.boot.u <- numeric(R)
  for(j in 1:R){
    Qs50.boot.l[j] <- x.c[which.min(abs(cum.boot.l[,j] - 0.5))]
    Qs50.boot.m[j] <- x.c[which.min(abs(cum.boot.m[,j] - 0.5))]
    Qs50.boot.u[j] <- x.c[which.min(abs(cum.boot.u[,j] - 0.5))]
  }
  
  cum.l <- apply(cum.boot.l, 1, quantile, probs = 0.975)
  cum.m <- apply(cum.boot.m, 1, quantile, probs = 0.5)
  cum.u <- apply(cum.boot.u, 1, quantile, probs = 0.025)

  Qeff.l <- as.numeric(quantile(Qeff.boot.l, probs = 0.025))
  Qeff.m <- as.numeric(quantile(Qeff.boot.m, probs = 0.5))
  Qeff.u <- as.numeric(quantile(Qeff.boot.u, probs = 0.975))
  
  spread.qeff <- abs(round((Qeff.u - Qeff.l) / Qeff.m,2))
  
  # Calc Qs50 and 95% CI
  Qs50.l <- round(as.numeric(quantile(Qs50.boot.l, probs = 0.025)),2)
  Qs50.m <- round(as.numeric(quantile(Qs50.boot.m, probs = 0.5)),2)
  Qs50.u <- round(as.numeric(quantile(Qs50.boot.u, probs = 0.975)),2)
  spread.qh <- round((Qs50.u - Qs50.l) / Qs50.m,2)
  
  #   plot(x.c, cum[,1], type = "l")
  #   lines(x.c, cum.l)
  #   lines(x.c, cum.u)
  #   abline(h=0.5)
  #   abline(v=Qs50.u)
  #   abline(v=Qs50.l)
  #   abline(v=Qs50.m)
  
  ## Write results to table  ADD SPREAD FOR QS50 AND QEFF HERE AS WELL
  results[i,] <- c(flow.sites$site.no[i], flow.sites$site.nm[i], 
                   n, q.min.sed, q.max.sed,
                   signif(b0.ci[1],3), b0, signif(b0.ci[3],3),
                   round(b1.ci[1],3), b1, round(b1.ci[3],3),
                   bcf.qmle, bcf.mvue, r2, sigma.e,
                   rl, n.eff, signif(h,3), signif(max.var.kdf,3),
                   q.min.flow, q.max.flow, cv, flash.RB,
                   Qs50.l, Qs50.m, Qs50.u, signif(spread.qh,3),
                   round(Qeff.l,2), round(Qeff.m,2), round(Qeff.u,2),
                   signif(spread.qeff,3))   
}

setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty")
write.csv(results, file = "resultsUncBootBed.csv", row.names = F)


############ Plots ##################

# Plot of four representations of propagated uncertainty in MFA
par(pty = "s", mfrow = c(2,2), mar = c(5,5,1,1), oma = c(1,1,1,1), cex.lab = 1.2, cex.axis = 1)

## Rating Curve ##
plot(newdata$cms, newdata$Qs_MVUE, type = "n", log = "xy", 
     xlab = expression("Discharge ("*m^3/s*")"), ylab = "Sediment Load (kg/s)", 
     #main = "Bedload Data Example - OLS",
     #      xlim = c(10^floor(log10(min(sed.site$cms))),
     #               max(sed.site$cms)),
     #      ylim = c(10^floor(log10(min(sed.site$kg.sec))),
     #               max(sed.site$kg.sec)),
)
col1 <- rgb(178/255,34/255,34/255, alpha = 0.5)
xx <- c(newdata$cms, sort(newdata$cms, decreasing = T))
yy <- c(newdata$UCL, sort(newdata$LCL, decreasing = T))
polygon(x = xx, y = yy, 
        col = col1, 
        border = NA)

lines(newdata$cms, newdata$Qs_MVUE, lwd = 1.5)
# lines(newdata$cms, newdata$LCL, col="red", lty=2) 
# lines(newdata$cms, newdata$UCL, col="red", lty=2)
points(sed.site$cms, sed.site$kg.sec)

# usr0 <- par("usr")       # get user coordinates
# par(usr = c(0, 1, 0, 1))
text(.11, 11, labels = "(a)", cex = 1)
# par(usr = usr0)

## Yield Record Example ##
plot(220:300,ucb[220:300], type = "n", 
     ylab = "Sediment Load (tonnes/day)", 
     xlab = "Days of Water Year")
xx <- c(220:300,300:220)
yy <- c(ucb[220:300], lcb[300:220])

col1 <- rgb(178/255,34/255,34/255, alpha = 0.5)
polygon(x = xx, y = yy, 
        col = col1, 
        border = NA)
lines(220:300,yield.mean[220:300], lwd=1.5)

usr0 <- par("usr")       # get user coordinates
par(usr = c(0, 1, 0, 1))
text(.075, .95, labels = "(b)", cex = 1)
par(usr = usr0)

## Yield PDF ##
plot(q.e, epdf[3,], type = "n",
     ylab = "Yield Probability Density",
     xlab = expression("Discharge"~~"("*m^3/s*")")
)
xx = c(q.e, sort(q.e, decreasing = T)) 
yy = c(epdf[3,], epdf[1,][length(epdf[1,]):1])

col1 <- rgb(65/255,105/255,225/255, alpha = 0.5)
col2 <- rgb(0,0,139/255, alpha = 0.5)
polygon(x = xx, y = yy, 
        col = col1, 
        border = NA,
        lwd=3)
# polygon(x = c(Qs50.l,Qs50.l,Qs50.u,Qs50.u), y = c(-10,0.5,0.5,-10),
#         border = NA,
#         col = col2)
lines(q.e, epdf[2,], lwd = 1.5)
segments(q.e[which.max(epdf[3,])], -10, q.e[which.max(epdf[3,])], max(epdf[3,]), lwd = 1.5, col = "gray30", lend = 1, lty = 2)
text(x = q.e[which.max(epdf[3,])]+1, y = 0, labels = expression(italic(Q[eff])), adj = c(0,0), cex = 1)
usr0 <- par("usr")       # get user coordinates
par(usr = c(0, 1, 0, 1))
text(.075, .95, labels = "(c)", cex = 1)
par(usr = usr0)

## Yield CDF ##
### Plot eCDF of Sediment Yield with Confidence Bands
par(pty = "s")
plot(flow.sort, yield.cum$ecdf.m, type = "n",
     ylab = "Yield Cumulative Probability",
     xlab = expression("Discharge"~~"("*m^3/s*")")
)
xx = c(flow.sort, sort(flow, decreasing = T)) #sort(flow, decreasing = T)
yy = c(l.ecdf, sort(u.ecdf, decreasing = T))#sort(flow, decreasing = T)

col1 <- rgb(65/255,105/255,225/255, alpha = 0.5)
col2 <- rgb(.5,.5,.5, alpha = 0.5)
polygon(x = xx, y = yy, 
        col = col1, 
        border = NA)
polygon(x = c(Qs50.l,Qs50.l,Qs50.u,Qs50.u), y = c(-10,0.5,0.5,-10),
        border = NA,
        col = col2)
lines(flow.sort, yield.cum$ecdf.m, lwd = 1.5)
segments(-10, 0.5, Qs50.u, 0.5, lwd = 1.5, col = "gray30", lend = 1, lty = 2)
segments(Qs50.m, -10, Qs50.m, 0.5, lwd = 1.5, col = "gray30", lend = 1, lty = 2)
text(x = 0, y = 0.52, labels = "Half Yield", adj = c(0,0), cex = 1)
text(x = Qs50.u+1, y = 0, labels = expression(italic(Q[s50])), adj = c(0,0), cex = 1)

usr0 <- par("usr")       # get user coordinates
par(usr = c(0, 1, 0, 1))
text(.075, .95, labels = "(d)", cex = 1)
par(usr = usr0)