####################################################################################
# MFA Uncertainty Analysis
# Joel Sholtes, Colorado State University
# December, 2014
#
#
# Algorithm for derived distribution Confidence Bands:
# a) Calc confidence and prediction Bands for sediment rating curve using OLS
# b) Calc Confidence band for flow distribution using Kernel Density / Normal Approx
# c) Calc upper, mean, and lower 95% CB/PB for sediment yield PDF / CDF using
#    empirical PDF method and numerical integration to generate CDF
# d) Uncertainty interval for Qeff and Qs50 is then the range between upper band for upper 
#    sed yield distribution lower band for lower sed yield distribution. 
####################################################################################

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

rm(list = ls())
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty")
# path <- "bedload_sites.csv"
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

# 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
  
  #   # remove zero flow values (need a way to deal with this...)
  #   flow <- flow[which(flow > 0)]
  
  # 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, ln.q = log(x), ln.qs = log(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(ln.qs ~ ln.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(ln.q = seq(log(min.flow)*1.1, max(log(flow))*1.1, 0.005))
            }else{
              data.frame(ln.q = seq(log(min.flow)*.9, max(log(flow))*1.1, 0.005))
            }
  y.conf <- predict(model.ols, newdata = x.pred, interval = "confidence") 
  y.pred <- predict(model.ols, x.pred, interval = "prediction") 
    
  ## 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])
        
###########################################################
  # Calculate parameters for re-transformation
  lnQ_bar <- mean(x)
  m = n - 2 # degrees of freedom
  # variance of discharges in regression dataset
  Sxx = sum((sed.site$ln.q - lnQ_bar)^2) 
    
  # Add vector of Q
  x.pred <- data.frame(x.pred, cms = exp(x.pred$ln.q))
    
  # Calculate MVUE BCF ->  see Cohn et al. (1989)
  # Variance of new compared with regression dataset Q
  var = 1/n + (x.pred$ln.q - lnQ_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
    
  # Area under ePDF
  area <- trapz(kde$x, kde$y)
    
############# 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  
    
############ Create Sediment Yield (effectiveness) Curve with CBs ##############
  # 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
  
  # Conf and Pred Bands for Qs relation in arithmatic space
  qs.cb <- rbind(exp(y.conf[,2]) * bcf, # Lower
                 exp(y.conf[,1]) * bcf, # Mean
                 exp(y.conf[,3]) * bcf  # Upper
                 )
  qs.pb <- rbind(exp(y.pred[,2]) * bcf, # Lower
                 exp(y.pred[,1]) * bcf, # Mean
                 exp(y.pred[,3]) * bcf  # Upper
                 )
  
  # Conf and Pred Bands for Eff Curve
  eff <- data.frame(mean = qs.cb[2,]*kde.cb2[2,],
                    low.conf = qs.cb[1,]*kde.cb2[1,],
                    upp.conf = qs.cb[3,]*kde.cb2[3,],
                    low.pred = qs.pb[1,]*kde.cb2[1,],
                    upp.pred = qs.pb[3,]*kde.cb2[3,]
                    )
  
  # Calculate cumulative integral of eff for cum sed yield curves
  x.c <- rollmean(x.pred$cms, k = 2)
  cum <- apply(eff, MARGIN = 2, FUN = cum.int, x = x.pred$cms)
  cum.u <- apply(cum[,2:5], MARGIN = 1, FUN = min)
  cum.l <- apply(cum[,2:5], MARGIN = 1, FUN = max)
  
  # Calc Qs50 and 95% CI
  Qeff.l <- approx(x = eff$low.pred, y = x.pred$cms, xout = max(eff$low.pred))$y
  Qeff.m <- approx(x = eff$mean, y = x.pred$cms, xout = max(eff$mean))$y
  Qeff.u <- approx(x = eff$upp.pred, y = x.pred$cms, xout = max(eff$upp.pred))$y
  spread.qeff <- abs(round((Qeff.u - Qeff.l) / Qeff.m,2))
  
  # Calc Qs50 and 95% CI
  Qs50.l <- approx(cum.l, x.c, xout = 0.5)$y
  Qs50.m <- approx(cum[,1], x.c, xout = 0.5)$y
  Qs50.u <- approx(cum.u, x.c, xout = 0.5)$y
  spread.qh <- round((Qs50.u - Qs50.l) / Qs50.m,2)
  
  ## 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, spread.qh,
                   Qeff.l, Qeff.m, Qeff.u, spread.qeff)     
}

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