tea.eo.plots: TEA's plotting functions

#' @title Stochastic Plot
#'
#' @description Plot stochastic results. PlotStochastic takes a long table, and PlotMatStochastic takes short table or matrix.
#'
#' @param x x An array of x variable.
#' @param ys An array or matrix of y variables.
#' @param col.f Color function with transparency index.
#' @param q1 An array of quantiles. Default set as c(0.05,0.95).
#' @param q2 An array of quantiles. Default set as c(0.25, 0.75).
#' @param add TRUE/FALSE 
#' @param ylim 
#' @param na.rm
#'
#' @export PlotStochastic
#' @export PlotMatStochastic
#' @export Plot2Stochastics
#' @export Plot2MatStochastics
#'
#' @examples
#' # Generate a long table
#' niter <- 100
#' df <- data.frame("Date"=rep(seq(as.Date("2010-1-1"), by="month", length.out = 12),niter)
#' , "Iter" = rep(1:niter,each=12), "Y"=rnorm(12*niter))
#' for (i in 2:nrow(df))if(diff(df$Iter)[i-1]==0) df$Y[i] <- df$Y[i] + df$Y[i-1]
#' PlotStochastic(df$Date, df$Y, blue.f, q1 = c(0.05,0.95), q2=c(0.25, 0.75), ylab="y")
#'
#' # 2 stochastic data in long table format
#' Plot2Stochastics(df$Date, df$Y, df$Y/2+2, red.f, blue.f, qs = c(0.05,0.95), ylab="y")
#'
#' # A short table
#' df <- cast(df, Date ~ Iter, value="Y")
#' PlotMatStochastic(df$Date, df[,-1], blue.f, q1 = c(0.05,0.95), q2=c(0.25, 0.75), ylab="y")


PlotStochastic <- function(x, ys, col.f=blue.f, q1 = c(0.05,0.95), q2=c(0.25, 0.75), add=FALSE, ylim=NULL, ...){

  if (is.null(ylim)) ylim=range(ys)
  if(!add) plot0(x, ys, col=col.f(0), ylim=ylim, ...)

  m <- aggregate(ys, by=list(x), mean, na.rm=TRUE)
  lines(m[,1], m[,2], col=col.f(1), lwd=3)

  # min/max
  temp.df <-aggregate(ys, by=list(x), quantile, c(0, 1), na.rm=TRUE)
  polygon(c(temp.df[,1], rev(temp.df[,1])), c(temp.df[,2][,1], rev(temp.df[,2][,2])), col=col.f(0), border=col.f(0.2))

  # quantile set 1
  temp.df <-aggregate(ys, by=list(x), quantile, q1, na.rm=TRUE)
  polygon(c(temp.df[,1], rev(temp.df[,1])), c(temp.df[,2][,1], rev(temp.df[,2][,2])), col=col.f(0.3), border=col.f(0.5))

  # quantile set 2
  temp.df <-aggregate(ys, by=list(x), quantile, q2, na.rm=TRUE)
  polygon(c(temp.df[,1], rev(temp.df[,1])), c(temp.df[,2][,1], rev(temp.df[,2][,2])), col=col.f(0.3), border=col.f(0.5))

  #DrawHLines(range(ys))

}

#' @describeIn PlotStochastic One Stochastic data series in a short table format
PlotMatStochastic <- function(x, ys, col.f=blue.f, q1 = c(0.05,0.95), q2=c(0.25, 0.75), ylim=NULL, add=FALSE, na.rm=FALSE, ...){

  if (na.rm){
    filtr <- !is.na(apply(ys,1, sum))
    x <- x[filtr]
    ys <- ys[filtr,]
  } 
  
  if (is.null(ylim)) ylim=range(ys)
  if(!add) plot0(x, ys[,1], col=col.f(0), ylim=ylim, ...)
  lines(x, apply(ys,1,mean), col=col.f(1), lwd=3)

  # min/max
  qs <-apply(ys,1, quantile, c(0, 1))
  polygon(c(x, rev(x)), c(qs[1,], rev(qs[2,])), col=col.f(0), border=col.f(0.2))

  # quantile set 1
  qs <-apply(ys,1, quantile, q1)
  polygon(c(x, rev(x)), c(qs[1,], rev(qs[2,])), col=col.f(0.3), border=col.f(0.5))

  # quantile set 2
  qs <-apply(ys,1, quantile, q2)
  polygon(c(x, rev(x)), c(qs[1,], rev(qs[2,])), col=col.f(0.3), border=col.f(0.5))

  #DrawHLines(range(ys))

}

#' @describeIn PlotStochastic Two stochastic data series in a long table format.
Plot2Stochastics <- function(x, y1, y2, col1.f, col2.f, qs = c(0.05,0.95), ylim=NULL, ...){

  if (is.null(ylim)) ylim=range(y1,y2)
  plot0(x, y1, col=col1.f(0), ylim=ylim, ...)

  # first y variable
  m <- aggregate(y1, by=list(x), mean, na.rm=TRUE)
  lines(m[,1], m[,2], col=col1.f(1), lwd=3)
  temp.df <-aggregate(y1, by=list(x), quantile, qs, na.rm=TRUE)
  polygon(c(temp.df[,1], rev(temp.df[,1])), c(temp.df[,2][,1], rev(temp.df[,2][,2])), col=col1.f(0.3), border=col1.f(0.5))

  # second y variable
  m <- aggregate(y2, by=list(x), mean, na.rm=TRUE)
  lines(m[,1], m[,2], col=col2.f(1), lwd=3)
  temp.df <-aggregate(y2, by=list(x), quantile, qs, na.rm=TRUE)
  polygon(c(temp.df[,1], rev(temp.df[,1])), c(temp.df[,2][,1], rev(temp.df[,2][,2])), col=col2.f(0.3), border=col2.f(0.5))

  #DrawHLines(range(y1, y2))

}

#' @describeIn PlotStochastic Two stochastic data series in a matrix format.
Plot2MatStochastics <- function(x, ys1, ys2, col1.f, col2.f, qs = c(0.05,0.95), ylim=NULL, ...){
  
  if (is.null(ylim)) ylim=range(ys1, ys2, na.rm=TRUE)
  plot0(x, ys1[,1], col=col1.f(0), ylim=ylim, ...)

  # first y variable
  filtr <- !is.na(apply(ys1,1, sum))
  lines(x[filtr], apply(ys1[filtr,], 1, mean), col=col1.f(1), lwd=3)
  qs.v <- apply(ys1[filtr,], 1, quantile, qs)
  polygon(c(x[filtr], rev(x[filtr])), c(qs.v[1,], rev(qs.v[2,])), col=col1.f(0.3), border=col1.f(0.5))

  # second y variable
  filtr <- !is.na(apply(ys2, 1, sum))
  lines(x[filtr], apply(ys2[filtr,], 1, mean), col=col2.f(1), lwd=3)
  qs.v <-apply(ys2[filtr,], 1, quantile, qs)
  polygon(c(x[filtr], rev(x[filtr])), c(qs.v[1,], rev(qs.v[2,])), col=col2.f(0.3), border=col2.f(0.5))

}