#######################################################################################
##  Log-linear Regression of Sediment Load-Discharge Data
##
##  Joel Sholtes, PhD Candidate 
##  Colorado State University, Dept. Civil and Env. Engr.
##  jsholte@gmail.com
##  May, 2014
##
##  Uses OLS and RLM (robust linear) regression techniques to fit a regression line
##  to the sediment load (tonnes/day) discharge (m^3/s) bivariate data.
##  Also calculates various correlation values, as well as bias correction factors.
##  Bias correction factors (BCFs) used include the following as outlined in 
##  Cohn, T (1995), "Recent Advances in Statistical Methods for the Estimation of Sediment 
##  and Nutrient Transport in Rivers", Reviews of Geophysics, 1117-1123. 
##  The two BCFs listed below are simply multiplication factors in the re-transformed sediment
##  rating curve: Qs = a * Q^b * BCF. Cohn et al. (1989) presente the
##  BCF_MVUE in the case of load estimation using a regression equation and a flow record.
##  It is currently unclear to me how to apply this to a sediment rating curve only.
##
##  BCF_QMLE = exp(s_e^2 / 2) e.g., Fergusen 1986 
##  BCF_SM = sum( exp(e_i^2) ) / N (e_i = individual residual error) e.g., Duan 1983
#######################################################################################

rm(list=ls(all=TRUE))

library(MASS)
library(boot)
library(robust)

setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Bedload")
#load Finney (1941) function
# source("Finney.R")

#######################################################################################
## Start from here (if data already imported and cleaned):
#######################################################################################

#compiled bedload data
sed.data <- read.table("BedLoad_R_Data.csv", stringsAsFactors = F,header = TRUE, sep = ",")

# most up-to-date list of sed sites
sed.sites <- read.table("site_summary.csv", stringsAsFactors = F, header = TRUE, sep = ",")
sites <- sed.sites$site.no

# Remove zero discharge load observations for log-transformation
sed.data <- subset(sed.data, !sed.data$cms == 0)

# # remove culled sites due to regulation and poor fit
# rm.sites <- read.table("List of Sites to Remove.txt", stringsAsFactors = F, header = T)
# data <- subset(data, !data$SITE_NO %in% rm.sites$site.id)

# # extracts unique site.id's
# sites <- unique(sed.data$site.id)
# sites <- subset(sites, sites %in% sed.sites$site.no)
# sites <- sort(as.numeric(sites))


# Create empty data frame for results and assign col names to it
cnames <- c("site.no", "site.nm", "n", "b0_RLM", "b1_RLM", "BCF_smear_RLM", "BCF_qmle_RLM",
            "b0_OLS", "b1_OLS", "BCF_qmle_OLS",
            "cor_S", "cor_K", "r2_OLS", "r2_RLM", "pb1_OLS", "pb1_RLM", "b1_ROB",
            "min.flow", "max.flow", "Q0")
results <- data.frame(matrix(data = 0, nrow = length(sites), ncol = length(cnames)),
                      row.names = NULL, stringsAsFactors = F)
names(results) <- cnames  # Assigns col names to data frame

# read in site numbers and calc rating curve, CI, and values of coeffs
# seems to work pretty well, but given rows of NAs.  Want to make sure it's
# b/c of missing data or too little data. Need to also think about adding in 
# missing data or flow values, which might exist in NWIS

# Also, need to add CI, and linear regression line to output tables
# as well as diagnostic plots?
for (i in 1:length(sites)){
  site_data <- subset(sed.data, site.id == sites[i])
  x <- numeric(nrow(site_data))
  y <- numeric(nrow(site_data))
  N <- nrow(site_data)
  
  # Extract Site specific flow - concentration pairs and create temp  
  # x = log(Q[cms]), y = log(Qs[tonnes/day]) vectors for regression
  if (nrow(site_data) >= 15){
    
    Q0 <- subset(site_data, site_data$kg.sec == 0)
    Q0 <- mean(Q0$cms)
    
    # Remove zero sed load observations for log-transformation
    site_data <- subset(site_data, !site_data$kg.sec == 0)
    
    x <- log(site_data$cms)
    y <- log(site_data$kg.sec)
    
    # Calc correlation using Spearman and Kendall coefficients
    # see this page for explanation of Pearson v. Spearman v. Kendall cor coeff:
    # http://www.unesco.org/webworld/idams/advguide/Chapt4_2.htm 
    # Spearman seems to be more similar to Pearson (nonparametric) and describes variance  
    # explained by the predictor variable
    corS <- signif(cor(x, y, use = "complete.obs", method = "spearman"), 2)
    corK <- signif(cor(x, y, use = "complete.obs", method = "kendall"), 2)
    
    # OLS Log-Linear Regression
    model_OLS <- lm(y ~ x)
    b0_OLS <- signif(exp(coef(summary(model_OLS))["(Intercept)","Estimate"]), 3)
    b1_OLS <- signif(coef(summary(model_OLS))["x","Estimate"], 3)
    BCF_qmle_OLS <- signif(exp(sum(residuals(model_OLS)^2)/(N-2)/2), 3)
    r2 <- signif(summary(model_OLS)$r.squared, 2)
    pb1_OLS <- signif(summary(model_OLS)$coefficients[2,4],3)
    
    # Robust Log-Linear regression
    model_RLM <- rlm(y ~ x, method = "MM", maxit = 100)
    b0_RLM <- signif(exp(coef(summary(model_RLM))["(Intercept)","Value"]), 3)
    b1_RLM <- signif(coef(summary(model_RLM))["x","Value"], 3)
    BCF_smear_RLM <- signif(sum(exp(residuals(model_RLM)))/N, 3) # produces some really large #'s sometimes
    BCF_qmle_RLM <- signif(exp(sum(residuals(model_RLM)^2)/(N-2)/2), 3)
    
    # Robust Log-Linear regression using other method that produces approx. p-vals
    rob.control <- lmRob.control(mxr = 100, mxf = 100, mxs = 100,)
    model_rob <- lmRob(y ~ x, control = rob.control)
    pb1_RLM <- signif(summary(model_rob)$coefficients[2,4],3)
    r2_RLM <- signif(summary(model_rob)$r.squared,2)
    b1_ROB <- signif(summary(model_rob)$coefficients[2,1],3)
    
    # Also want to include CI's for coeffs 
    
    # Write data to table along with CI for regression line, which may be of different length?
    data_pairs <- data.frame(cms = exp(x), kg.sec = exp(y))
    write.csv(data_pairs, paste("sediment/" , sites[i], ".csv", sep = ""), row.names = F) # Need to make this folder first
    
    results[i,] <- c(sites[i], sed.sites$site.nm[i], N, 
                     b0_RLM, b1_RLM, BCF_smear_RLM, BCF_qmle_RLM, 
                     b0_OLS, b1_OLS, BCF_qmle_OLS,
                     corS, corK, r2, r2_RLM, pb1_OLS, pb1_RLM, b1_ROB,
                     signif(min(exp(x)), 3), round(max(exp(x)), 2),
                     signif(Q0, 3)
                     )
  }
  else {
    # Write only SITE_ID and N to table if N < 15
    results[i,1] <- sites[i]
    results[i,2] <- site_data$STATION_NM[1]
    results[i,3] <- N
  }
}

rm(b0_OLS, b0_RLM, b1_OLS, b1_RLM, b1_ROB, BCF_qmle_OLS, BCF_qmle_RLM,
   BCF_smear_RLM, pb1_OLS, pb1_RLM, Q0, r2_RLM)
attach(results)
boxplot(as.numeric(b1_OLS), at = 1, xlim = c(0.5, 3.5), ylim = c(0,6), 
        width = 2,
        cex.lab = 1.2,
        ylab = "Beta")
boxplot(as.numeric(b1_RLM), at = 2, width = 2, add = T)
boxplot(as.numeric(b1_ROB), at = 3, add = T, width = 2)
axis(1, at = 1:3, labels = c("OLS", "RLM", "ROB"))
detach(results)

write.csv(results,"results_bed_regression_ROB.csv", row.names = F)


#######################################################################################
## Analysis of Results
#######################################################################################
# Import previous results again
rm(list=ls(all=TRUE))
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Bedload")
results <- read.csv("results_bed_regression.csv", stringsAsFactors = F)
# remove sites with N < 15 for subsequent analysis 
results <- subset(results, !results$b0_RLM == 0)
# sort by siteID
results <- results[order(results$siteID), ]

attach(results)
# Analyze beta values from different methods
pdf( "b_OLS-b_RLM.pdf", width = 4, height = 4)
par(mfrow = c(1,1), pty = "s")
plot(b1_RLM, b1_OLS, xlab = expression(italic(beta) [RLM]), 
     ylab = expression(italic(beta) [OLS]), xlim = c(0,5), ylim = c(0,5))
abline(0,1)
dev.off()

par(mfrow = c(1,2))
hist(b1_RLM,15, xlab = expression(beta [RLM]), main = NULL)
hist(b1_OLS,15, main = NULL, xlab = expression(beta [OLS]))

par(mfrow = c(1,2))
qqnorm(b1_RLM, main = expression(beta[RLM]))
qqline(b1_RLM)
qqnorm(b1_OLS, main = expression(beta[OLS]))
qqline(b1_OLS)

library(vioplot)
par(mfrow = c(1,1))
pdf( "BetaVals.pdf", width = 4, height = 4)
vioplot(b1_RLM, b1_OLS, names = c("RLM", "OLS"), col = "gold")
title(ylab = bquote(beta), cex = 2)
dev.off()

mean(b1_RLM)
median(b1_RLM)
mean(b1_OLS)
median(b1_OLS)

t.test(b1_RLM, b1_OLS)
wilcox.test(b1_RLM, b1_OLS, paired = FALSE, exact = TRUE)

