############################################################################
## Record Extension Script Following Moog et al. 1999 (WRR)
##  "Streamflow record extension using power transformations"
##
## Joel Sholtes - Colorado State University, Fort Collins, CO
## April, 2014
## jsholtes@gmail.com
##
## Uses Box-Cox transformation of at site short-term record and 
## nearby reference (base) gage to create best linear relationship
## between the two. Uses MOVE.1 approach to maintain variance in 
## flow record (Hirsch 1982). Working with original data in CFS.
## does not work if there are zero flow days in either data set.
## Creates PDF of plots of analyses for each site used to extend a record.
############################################################################

rm(list = ls())
library(segmented) # Piece-wise regression
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Flow Record Extension")

## Import table of flow records to be extended
path.sites <- "extended_susp_sites.csv"
version <- 1
ext.sites <- read.table( path.sites, header = T, stringsAsFactors = F, sep = ",")

# This table contains slope and intercept values previously calculated using piecewise regression
# see lines 181 - 206, Must run code one site at a time and manuall yinput segmentation data
# First run code with empty seg.data table (slope1 = 0 for all rows), then decide if
# segmentation is necessary from initial output plots
seg.data <- read.table("ext_summary_segmented_susp.csv", header = T, stringsAsFactors = F, sep = ",")

# create data frame to summarize transformation results
cnames <- c("site.id", "site.nm", "site.start", "site.end", "site.yrs",
            "base.id", "base.nm", "base.start", "base.end", "base.yrs",
            "max.orig", "max.ext", "sd.orig", "sd.ext", "mean.orig", "mean.ext",
            "R^2", "SSE", "lambda.x", "lambda.y", "lx-ly", "complex")

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

# select flow extensions to perform 
rows <- 1:nrow(ext.sites)
rows <- c(1,8)

# OR load existing ext.summary table if re-running some sites one-at-a-time
#ext.summary <- read.table("ext_summary8.csv", header = T, stringsAsFactors = F, sep = ",")

#####################################################################################
# Loop to calc all optimal BC transformations for each base station and site station
# and create extended flow record based on the regression
#####################################################################################

# Optional PDF of four figures per site: 1) Concurrent site and base flow records,
# 2) Extended and original flow record superimposed, 
# 3) SSE/(n+1) of regression surface as a fucntion of lambda.x and lambda.y, and 
# 4) Piecewise and single regression lines of concurrent flow data on transformed axes

pdf("Flow_Ext_Figs_Susp_OLD.PDF", width = 8.5, height = 11)
  par(mfrow = c(2,2), pty = "s", mar = c(4,4,2,2), oma = c(3,3,3,3))
  
  for (i in rows){
    ######################################################################################
    # Import flow data 
    ######################################################################################
    
    # import flow.site station record and extended flow.base record
    # "site data" (short record)
    site.id <- ext.sites$id.site[i]
    path <- paste("Susp_Ext_Site/Formatted/", site.id, ".csv", sep = "")
    site.data.pre <- read.table(path, header = T, stringsAsFactors = F, sep = ",")
    # This turns both table formats into a table with only two columns: date, val (cfs)
    if (ncol(site.data.pre) == 2 ) { # two different site data types, some with two or four columns
        site.data <- data.frame(date = site.data.pre$date, val = site.data.pre$val)
    }else{
        site.data <- data.frame(date = site.data.pre$dates, val = site.data.pre$val)
    }
    # Format date character to be date format in R
    site.data$date <- as.Date(site.data$date, format = "%Y-%m-%d", origin = "1970-01-01")
    
    # Extract start date and end date from site data 
    # (might be a larger time span than the base station record)
    start.date.site <- as.Date(site.data$date[1], format = "%Y-%m-%d", origin = "1970-01-01")
    end.date.site <- as.Date(site.data$date[nrow(site.data)], format = "%Y-%m-%d", origin = "1970-01-01")
    site.yrs <- nrow(site.data)/365
    
  #   # add a day or two (lag) to the site flow date based on ext.sites table
  #   lag <- ext.sites$lag[i]
  #   dates <- as.numeric(site.data$date) + lag
  #   site.data$date <- as.Date(dates, format = "%Y-%m-%d", origin = "1970-01-01")
    
    # "base station" data used to extend the flow record
    base.id <- ext.sites$id.base[i]
    path <- paste("Susp_Ext_Base/", base.id, ".csv", sep = "")
    base.data <- read.table(path, header = T, stringsAsFactors = F, sep = ",")
    base.data$date <- as.Date(base.data$date, format = "%Y-%m-%d", origin = "1970-01-01")
    
    start.date.base <- as.Date(base.data$date[1], format = "%Y-%m-%d", origin = "1970-01-01")
    end.date.base <- as.Date(base.data$date[nrow(base.data)], format = "%Y-%m-%d", origin = "1970-01-01")
    base.yrs <- nrow(base.data)/365
    
    # Grab first and last dates on longest flow record to create continuous flow record.
    # This is important where there are gaps in base flow record but not in site flow record.
    # Or, the site flow record may actually be longer than the base flow record
    start.date <- as.Date(ifelse(start.date.base > start.date.site, start.date.site, start.date.base),
                          origin = "1970-01-01")
    end.date <- as.Date(ifelse(end.date.site > end.date.base, end.date.site, end.date.base),
                                        origin = "1970-01-01")
    # This vector of dates is used later on to create an extended flow record
    dates.ext <- seq(from = start.date, to = end.date, by = "day")
    
    # base station concurrent records
    # subset the base station data so that they are concurrent with site data
    base.data.sub <- subset(base.data, base.data$date %in% site.data$date)
    xc <- base.data.sub$val
    xc[which(xc == 0)] <- 0.0001
    
    # Site station:
    # subset the site station to make sure it matches the length of the concurrent
    # base station record
    site.data.sub <- subset(site.data, site.data$date %in% base.data$date)
    yc <- site.data.sub$val
    n <- length(yc)
    yc[which(yc == 0)] <- 0.0001
    
    #################################
    # Optimal Box-Cox Transformation
    #################################
    
    # define Box-Cox transformation values as per Moog et al. (1999)
    #lambda <- c(-0.67, -0.5, -0.33, -0.25, -0.10, 0, 0.1, 0.25, 0.33, 0.5, 0.67)
    lambda <- round(seq(-1, 1, length = 61), 3)
    
    # create matrices of xb & yb:
    # The Box-Cox tranformed flow values (columns) for each values of lambda (rows)
    xb <- matrix(0, length(lambda), length(yc))
    yb <- matrix(0, length(lambda), length(yc))
    
    # nu.x & nu.y are the normalized BC-transformed flow values following Moog et al. (1999)
    # cols: transformed, normalized flow values; rows: lambda values
    nu.x <- matrix(0, length(lambda), length(yc))
    nu.y <- matrix(0, length(lambda), length(yc))
    
    # transform data using all combinations of lambda and fill in matrices
    # if lambda = 0, then the BC transformation is log(x)
    for (j in 1:length(lambda)) {
      xb[j,] <- if( lambda[j] == 0 ){ log(xc) }else{ (xc^lambda[j] - 1 ) / lambda[j] }
      nu.x[j,] <- xb[j,] / ( exp(mean(log(xc))) ^ (lambda[j] - 1) )
      
      yb[j,] <- if( lambda[j] == 0 ){ log(yc) }else{ (yc^lambda[j] - 1 ) / lambda[j] }
      nu.y[j,] <- yb[j,] / ( exp(mean(log(yc))) ^ (lambda[j] - 1) )
      
    }
    
    # Create empty matrix and data frame to store SSE values
    bc.mat <- matrix(0, nrow = length(lambda), ncol = length(lambda))
    rownames(bc.mat) <- lambda
    colnames(bc.mat) <- lambda
    
    # Calc SSE for each combination of lambda
    # here we create a matrix and a data frame to find min SSE
    for (j in 1:length(lambda)) {
      for (k in 1:length(lambda)){
        fit <- lm(nu.y[j,] ~ nu.x[k,])
        SSE <- sum(fit$residuals^2)/(n + 1)
        bc.mat[j,k] <- SSE
      }
    }
    
    # Extract lambda Values
    lambda.idx <- which(bc.mat == min(bc.mat), arr.ind = TRUE)
    lambda.x <- round(lambda[lambda.idx[2]], 3)
    lambda.y <- round(lambda[lambda.idx[1]], 3)
    
    # Calc fit of data to transformation best fit lambda values
    xcp <- if( lambda.x == 0 ){ log(xc) }else{ (xc^lambda.x - 1 ) / lambda.x }
      
    ycp <- if( lambda.y == 0 ){ log(yc) }else{ (yc^lambda.y - 1 ) / lambda.y }
    
    # calc mulitple R^2 value for transformed data
    r2 <- round(cor(xcp, ycp, method = "pearson")^2, 2)
    
    # Calc SSE from linear regression of transformed concurrent flow data
    SSE.min <- signif(min(bc.mat), 3)
    
    # Calc slope and intercept of BC regression line using data transformed with chosen lambda values 
    slope.p <- sd(ycp)/sd(xcp)
    intercept.p <- mean(ycp)
  
    #################################################################
    # Piece-wise regression to account for over-estimating high flows
    #################################################################
    
    # Conduct segmented regression using best guess values of upper break point
    # NOTE: This is conducted site-by-site and not in the loop. 
    # Need to first find the slopes and breaks, add them to the seg.data table and use the 
    # seg.data table in the loop to calc the stats and extend the flow records (2 steps)
    # Need to repeat the segmented regression until a visually good piecewise regression results
  
    x.s <- xcp[which(xcp>1.7)]
    y.s <- ycp[which(xcp>1.7)]
    lm <- lm( y.s ~ x.s )
    seg.lm <- segmented(lm, seg.Z =~x.s, psi = list(x.s = 2.3))
    plot(xcp,ycp)
    plot.segmented(seg.lm, add = T, lwd = 2, col = "red", conf.level = 0.95, shade = F)
    int1 <- intercept(seg.lm)$x.s[1]
    int2 <- intercept(seg.lm)$x.s[2]
    slp1 <- slope(seg.lm)$x.s[1]
    slp2 <- slope(seg.lm)$x.s[2]
    break.x <- summary.segmented(seg.lm)$psi[2]
    break.flow <- round( if( lambda.x == 0 ){
                           ( lambda.y*( intercept.p + slope.p*(break.x - mean(xcp)) ) + 1 )^(1/lambda.y)
                           }else{ 
                            ( lambda.y*( intercept.p + slope.p*(break.x - mean(xcp)) ) + 1 )^(1/lambda.y)
                           }, 2)
    
    #################################################################
    ## Now transform long term record based on box-cox regression: 
    #################################################################
  
    # pull segmented piecewise regression coeffecients from seg.data table (previously calculated)
    a1 <- seg.data$int1[i]
    a2 <- seg.data$int2[i]
    b1 <- seg.data$slope1[i]
    b2 <- seg.data$slope2[i]
    
    x.break <- (a1 - a2) / (b2 - b1) # simultaneous solution of two lines
    y.break <- a1+ b1*x.break
  
    # 1) transform entire record based on BC regression from base station using segmented
    #    piecewise regression coefficients ( if they exist: seg.data$slope1[i] != 0 )
    if (seg.data$slope1[i] != 0) {
      # 1a) transform flow record based on lambda.x
      #     we need these values to differentiate the upper and lower segments of 
      #     the piecewise regression based on x.break
      ye.t <- if(lambda.x == 0) {
                log(base.data$val)
                } else{
                  (base.data$val^lambda.x - 1 ) / lambda.x
                }
      # 1b) un-transform base flow record for both segments of the piecewise regression
      ye.pre1 <- if(lambda.x == 0 & lambda.y == 0){
              exp(a1 + b1*log(base.data$val) )
            }else if(lambda.x == 0){
              ( lambda.y*( a1 + b1*log(base.data$val) ) + 1 )^(1/lambda.y)
            }else if(lambda.y == 0){
              exp( a1 + b1*( ( base.data$val^lambda.x - 1 ) / lambda.x ) )
            }else{
              ( lambda.y*( a1 + b1*( base.data$val^lambda.x - 1 ) / lambda.x ) + 1 )^(1/lambda.y)
            }
      
      ye.pre2 <- if(lambda.x == 0 & lambda.y == 0){
              exp(a2 + b2*log(base.data$val) )
            }else if(lambda.x == 0){
              ( lambda.y*( a2 + b2*log(base.data$val) ) + 1 )^(1/lambda.y)
            }else if(lambda.y == 0){
              exp( a2 + b2*( ( base.data$val^lambda.x - 1 ) / lambda.x ) )
            }else{
              ( lambda.y*( a2 + b2*( base.data$val^lambda.x - 1 ) / lambda.x ) + 1 )^(1/lambda.y)
            }
      
      # 1c) For ye.t < x.break, this populates ye.pre with ye.pre1, otherwise it populates with ye.pre2
      # ifelse() matches the indices for each T and F
      # for example: ifelse(c(T,F,F,T), c(1,2,3,4), c(5,6,7,8)) = 1 6 7 4
      ye.pre <- ifelse(ye.t < x.break, ye.pre1, ye.pre2)
      
      } else{ # for no segmented relationship (or values) in the seg.data table
        ye.pre <- if(lambda.x == 0 & lambda.y == 0){
                    exp( intercept.p + slope.p*( log(base.data$val) - mean(xcp) ) )
                  }else if(lambda.x == 0){
                    ( lambda.y*( intercept.p + slope.p*( log(base.data$val) - mean(xcp) ) ) + 1 )^(1/lambda.y)
                  }else if(lambda.y == 0){
                    exp( intercept.p + slope.p*( ( base.data$val^lambda.x - 1 ) / lambda.x - mean(xcp) ) )
                  }else{
                    ( lambda.y*( intercept.p + slope.p*( (base.data$val^lambda.x - 1)
  				  /lambda.x - mean(xcp) ) ) + 1 )^(1/lambda.y)
                  }
      }
  
    # 2) populate the extended record (ye) with the continuous time series (dates) 
    #    and regressed extended record based on concurrent dates between base.data
    #    and the overall record of dates
    ye <- data.frame(date = dates.ext, val = NA)
    ye$val[ye$date %in% base.data$date] <- round(ye.pre, 2) 
    
    # 3) back substitute in the existing original site.data flow record
    ye$val[ye$date %in% site.data$date] <- round(site.data$val, 2)
    
    # determine if a singularity or complex number results from re-transformation
    # calc BC-transformed and regressed (Eqn. 4 in Moog et al. 1999) values of extended record: yep_hat
    # then find if lambda.y*yep_hat is <= 1 for all values
    yep_hat <- if( lambda.x == 0 ){
      intercept.p + slope.p*(log(base.data$val) - mean(xcp)) 
    }else{ 
      intercept.p + slope.p*( (base.data$val^lambda.x - 1 ) / lambda.x - mean(xcp) )
    }
    
    complex <- sum(lambda.y*yep_hat <= -1) >= 1
    
    # Calc properties of original and extended flow records
    max.orig <- max(site.data$val, na.rm = T)
    max.ext <- max(ye$val, na.rm = T)
    sd.orig <- round(sd(site.data$val, na.rm = T), 2)
    sd.ext <- round(sd(ye$val, na.rm = T), 2)
    mean.orig <- round(mean(site.data$val, na.rm = T), 2)
    mean.ext <- round(mean(ye$val, na.rm = T), 2)
    
    # Write flow extension results to ext.summary table
    ext.summary[i,] <- c(ext.sites$id.site[i], ext.sites$nm.site[i], format(start.date.site,"%Y-%m-%d"), 
                         format(end.date.site,"%Y-%m-%d"), site.yrs,
                         ext.sites$id.base[i], ext.sites$nm.base[i], format(start.date.base,"%Y-%m-%d"),
                         format(end.date.base,"%Y-%m-%d"), base.yrs,
                         max.orig, max.ext, sd.orig, sd.ext, mean.orig, mean.ext,
                         r2, SSE.min, lambda.x, lambda.y, abs(lambda.x - lambda.y), complex)
    
    # Save extended flow record
    path <- paste("Susp_Ext_Records/", ext.sites$id.site[i], ".csv", sep = "")
    write.table(ye, path, row.names = F, sep = ",")
    
    #################
    ## Plot Stuff
    #################
    
    ## Plot concurrent records
    plot(base.data.sub$date, base.data.sub$val, type = "n", xlab = "Year", ylab = "Flow [cfs]", 
         cex.axis = 0.75, cex.main = .75, main = "Concurrent Flow Records", log = "y", 
         ylim = c(10^(floor(log10(min(yc)+.001))), max( max(xc)*1.25, max(yc)*1.25 )))
    lines(site.data.sub$date, site.data.sub$val, lwd = 2, col = "dodgerblue")
    lines(base.data.sub$date, base.data.sub$val)
    par (font = 2)
    legend("bottomleft", # places a legend at the appropriate place 
           c("Site Gage Record", "Extension Gage Record"), # puts text in the legend 
           lty=c(1,1), # gives the legend appropriate symbols (lines)
           lwd=c(3,2),col=c("dodgerblue", "black"), # gives the legend lines the correct color and width
           cex = 0.75, # changed legend text size
           bty = "n")  # removes legend box)  
    
    ## Plot extended flow record alongside original
    title <- "Extended Flow Record"
    plot(ye$date, ye$val, type = "n", xlab = "Year", ylab = "Flow [cfs]", 
         cex.axis = 0.75, cex.main = 0.9, main = title)
    lines(site.data$date, site.data$val, lwd = 2, col = "dodgerblue")
    lines(ye$date, ye$val)
    
    ## Plot SSE as a function of lambda.x, lambda.y
    # create logarithmic contours
    ctr.min <- floor(log10(signif(min(bc.mat), 2))) # Calcs floor of log of min SSE
    ctr.max <- log10(ceiling(max(bc.mat))) # Calcs ceiling of log of max SSE
    ctrs <- signif( 10^( seq(ctr.min, ctr.max, length.out = 15) ), 2)
  
    # Transposing matrix here. I've double-checked that lambda.x and lambda.y are
    # indeed the pair that mimimizes SSE based on nu.y ~ nu.x, but for some reason
    # that I cannot figure out, when I plot that point on the contour plot, it comes
    # out switched.
    par(cex.main = 1, font = 1)
    title <- paste("SSE of Box-Cox Transformation \nSSE.min = ", SSE.min, sep = "")
    contour( x = lambda, y = lambda, z = t(bc.mat), main = title,
             xlab = bquote(lambda[x]), ylab = bquote(lambda[y]), 
             levels = ctrs, cex.axis = 0.75, cex.main = 0.9)
    points(lambda.x, lambda.y, pch = 19)
    
    ## Plot transformed data on transformed axes
    title <- paste("Regression with Transformed Data")
    plot(xcp, ycp, xlab = bquote(italic("x'")[c]), ylab = bquote(italic("y'")[c]),
         main = title,cex.axis = 0.75, cex.main = 0.9)
    curve(intercept.p + slope.p*(x - mean(xcp)), from = min(xcp), to = max(xcp), n = 2, 
          add = T, col = "blue", lwd = 3, lty = 2)
    if (seg.data$slope1[i] != 0) {
    curve(a1 + b1*x, from = min(xcp), to = x.break, n = 2, add = T, col = "red", lwd = 2)
    curve(a2 + b2*x, from = x.break, to = max(xcp), n = 2, add = T, col = "red", lwd = 2)
    }
  
    legend("topleft", # places a legend at the appropriate place 
           c("Segmented Reg.", "MOVE.1 Reg."), # puts text in the legend 
           lty=c(1,2), # gives the legend appropriate symbols (lines)
           lwd=c(2,3),col=c("red", "blue"), # gives the legend lines the correct color and width
           cex = 0.75, # changed legend text size
           bty = "n")  # removes legend box)  
    xcoord <- min(xcp) + ( max(xcp) - min(xcp) ) * 0.8
    ycoord1 <- min(ycp) + ( max(ycp) - min(ycp) ) * 0.11
    ycoord2 <- min(ycp) + ( max(ycp) - min(ycp) ) * 0.01
    ycoord3 <- min(ycp) + ( max(ycp) - min(ycp) ) * 0.21
    
    label1 <- bquote(lambda[x] == .(lambda.x))
    label2 <- bquote(lambda[y] == .(lambda.y))
    label3 <- bquote(R^2 == .(r2))
    
    text(x = xcoord, y = ycoord1, labels = label1, cex = .8, adj = c(0,0))
    text(x = xcoord, y = ycoord2, labels = label2, cex = .8, adj = c(0,0))
    text(x = xcoord, y = ycoord3, labels = label3, cex = .8, adj = c(0,0))
    
    title <- paste("Site: ", ext.sites$id.site[i], " - ", ext.sites$nm.site[i], 
                   "\nBase:", ext.sites$id.base[i], " - ", ext.sites$nm.base[i])
    mtext(title, line = -1, side = 3, outer = T)
  }
# Closes PDF program
dev.off()

file.name <- paste("ext_summary_susp_OLD", version, ".csv", sep = "")
write.table(ext.summary, file.name, row.names = FALSE, sep = ",")
