# -------------------------------------------------------------------------------------------------------------------------  #
  #  -----------------------------------------------------------------------------------------------------------------------   #
  #                                                                                                                            #
  #           File Name              :  beta plot.r                                                                            #
  #           Last Updated Funclist  :  15 Feb 2014,  1:13 AM (Saturday)                                                       #
  #                                                                                                                            #
  #           Author Name            :  Rick Saporta                                                                           #
  #           Author Email           :  RickSaporta@gmail.com                                                                  #
  #           Author URL             :  www.github.com/rsaporta                                                                #
  #                                                                                                                            #
  #           Packages Called        :  NA                                                                                     #
  #           Packages Used via NS   :  NA                                                                                     #
  #                                                                                                                            #
  #  -----------------------------------------------------------------------------------------------------------------------   #
  #                                                                                                                            #
  #   BetaPlot           ( a.seq, b.seq, x.min=0, x.max=1, pause=0.35, plot.this.many=length(a.seq), ci=0.9 )                  #
  #   find.ci.single.mode ( vec, ci=0.9, starting.val=NULL, starting.index=NULL                                                #
  #                         , use=c("index.length", "support.range", "mode"), support )                                        #
  #   checkInterval      ( ci.indxs, retval=TRUE )                                                                             #
  #   findClosestIndx    ( vec, value, direction=c("right", "left", "inward"), beyond=NA )                                     #
  #                                                                                                                            #
  #                                                                                                                            #
  #                                                     <END FUNCS>                                                            #
  #  -----------------------------------------------------------------------------------------------------------------------   #
  # -------------------------------------------------------------------------------------------------------------------------  #

#  TO CREATE RANDOM INPUT: 
# 
#    # iterate over several changing a, b for a fixed set of X
#    N <- 10
#    p <- 0.5
#    a.init <- 3
#    b.init <- 2
#    seed <- 7
#  
#    # calculate the series of updates to a, b, based on flips
#    change <- {set.seed(seed); rbinom(N, 1, p)} 
#    a.seq <- a.init + cumsum(change)
#    b.seq <- b.init + cumsum(abs(1-change))
#  



BetaPlot <- function(a.seq, b.seq, x.min=0, x.max=1, pause=0.35, plot.this.many=length(a.seq), ci=.9)  {

  x <- seq(from=x.min, to=x.max, length.out=250)

  # a & b should be the same length
  if (length(a.seq) != length(b.seq))
    stop("`a.seq` and `b.seq` must be the same length.")

  # caluclate Beta, then set names
  output <- mapply(Beta, a=a.seq, b=b.seq, MoreArgs=list(x=x), SIMPLIFY=FALSE)
  output <- setNames(output, paste0("a=", a.seq, ",  b=", b.seq))
  L <- length(output)

  # default to the length of the number of arguments, if 
  if (missing(plot.this.many) || !is.numeric(plot.this.many))
    plot.this.many <- L

  if (plot.this.many < L) {
    # if only 2, then select first and last
    if (plot.this.many==2)
      selected <- c(1, L)
    else 
      selected <- ceiling(c(1, seq(from=2, to=L, length.out=(plot.this.many-1))))
  
    # crop output and reset L. (note: output is already named)
    output <- output[selected]
    L <- length(output)
  }

# legend('topright', names(a)[-1] , 
#    lty=1, col=c('red', 'blue', 'green',' brown'), bty='n', cex=.75)


  # check for user error on the value of pause
  if (isTRUE(pause))
    pause <- 0.35
  if(!is.numeric(pause))
    pause <- 0
  if (pause < 0) {
    warning("pause should be a positive numeric value in seconds. Setting to 0.")
    pause <- 0
  }

## TODO: 
# set to just one for-loop.  First create the params, then change the values of the first and last param. 
#    then can easily create the legend 

  # plot initial value
  y.max <- max(unlist(output))
  plot(x=x, y=output[[1]], t="l", lwd=2, col=rgb(1, .10, .30, 1), ylim=c(0, y.max))


  # some values for iterating
  colors <- make.rgba(L-1, start=c(.8, .1, .7, .05), end=c(0, .8, 0, 1), color="log", final.alpha.gap=0.5)
  lwds <- seq(0.2, 1.5, length.out=L-1)

  # iterate
  for (i in 1:(L-1)) {
    Sys.sleep(pause)
    lines(x=x, y=output[[i+1]], t="l", lwd=lwds[[i]], lty=4, col=colors[[i]])
  }

  # add emphasis for the last plot
  lines(x=x, y=output[[i+1]], t="l", lwd=2, col=colors[[i]])

  ## find some stats
  out <- output[[L]]
  
  # mode
    mde.val  <- max(out) 
    mde.indx <- which(out == mde.val)
    mde.x.val <- x[[mde.indx]] 
    color.mde <- rgb(r=0, g=0.5, b=0.8, a=0.4)
    abline(v=mde.x.val, lwd=0.5, lty=8, col=color.mde)

  # ci bounds
    color.ci <- rgb(r=0, g=0.5, b=0.9, a=0.8)
    ci.lims <- find.ci.single.mode(out, ci=ci, use="mode")
    x.intv  <- x[ci.lims]
    abline(v=x[ci.lims], lwd=1, lty=2, col=color.ci)

    # horizontal bar
    lines(x=x[ci.lims], y=out[ci.lims], lwd=1.5, lty=2, col=color.ci)

    names(x.intv) <- names(ci.lims)
    x.intv <-  fw3(x.intv, digs=2)
    md <- paste0("Posterior Mode is at ", round(mde.x.val, 3))
    ci.msg <- paste0("A ", fwp(ci, 0, sep=""), " credible interval\n"
                    ,"on the posterior is ["
                    , x.intv[["lower"]], ", ", x.intv[["upper"]], "]")

  legend("topleft", c(names(output), md))
  title(paste("Beta distribution", ci.msg, sep="\n"))
  return(output[[L]])
}

find.ci.single.mode <- function(vec, ci=.9, starting.val=NULL, starting.index=NULL
                       , use=c("index.length", "support.range", "mode"), support) {

  use <- match.arg(use)

  # nesting this function to allow for variable access without having to pass the vars
  checkInterval <- function(ci.indxs, retval=TRUE) {
    intv <- (diff(ci.indxs) + 1) / length(vec)
    pass <- intv > ci
    if(!retval)
      return(pass)
    names(intv) <- as.character(pass)
    return(intv)
  }

  if (!is.null(starting.val) & !is.null(starting.index)) {
    warning("Both `starting.index` and `starting.val` have been supplied. Will disregard index and use only value.")
  }

  if (!is.null(starting.val)) {
    # find index based on starting.val
    mde.val <- starting.val
    mde <- ceiling(median(which(vec==mde.val)))
  } else if (!is.null(starting.index)) {
    # use starting.index
    mde <- starting.index
    mde.val <- out[mde]
  } else {
    # calculate the mode, and use that. 
    mde.val <- max(vec) 
    mde <- ceiling(median(which(vec==mde.val)))
  }

  ci.half <- setNames(c((1-ci)/2, 1-(1-ci)/2), c("lower", "upper"))

  # find ci.limits based on the index values of the current distribution
  if (use=="index.length")  {
    l.vec <- length(vec)
    ci.lims <- quantile(seq(l.vec), ci.half)

    # normally, round the index up, unless the bottom is already rounded down,
    #  in which case follow that direction. (Otherwise, the CI will overshoot slightly)
    if (as.logical(all.equal(ci.lims[[1]], floor(ci.lims[[1]]), tol=.0001)) )
      ci.indxs <- floor(ci.lims)
    else 
      ci.indxs <- ceiling(ci.lims)

  } else if(use=="support.range") {
    if (missing(support))
      stop("in order to use `support.range`, support must be passed explicitly.")
    s.range <- range(support, na.rm=TRUE)  
    ci.lims <- quantile(s.range, ci.half)

    ci.indxs <- findClosestIndx(support, ci.lims, direction="inward")
    checkInterval(ci.indxs)

  # find ci.lims based on mde
  } else if(use=="mode") {
    indxs.init <- c(mde, mde)
    max.number.of.guesses <- 1000
    speed.factor <- 3
    top <- max(vec)
    bot <- min(vec)

    # intial values
    if (speed.factor <= 1)
      stop("`speed.factor` must be greater than 1")
    indxs <- indxs.init
    counter <- 0
    guesses <- bot +  ((top-bot) * (1/speed.factor)^(1:max.number.of.guesses))


    while (!checkInterval(indxs, retval=FALSE) && counter < max.number.of.guesses) {
      counter <- counter+1
      guess <- guesses[[counter]] 
      indxs <- findClosestIndx(vec, value=guess, direction="inward", beyond=indxs)
      # checkInterval(indxs)
    }

    if (!checkInterval(indxs, retval=FALSE) )
      stop("Could not find CI in ", max.number.of.guesses, " tries. Try increasing the speed factor or the max.number.of.guesses.")
    # get finer
    else {
        prev.indxs <- indxs
        indxs.init <- indxs
        top <- guesses[max(1, counter-2) ]
        bot <- guess
        speed.factor <- 1.2
        max.number.of.guesses <- max.number.of.guesses
        indxs.init <- findClosestIndx(vec, value=top, direction="inward", beyond=mde)


        # intial values
        if (speed.factor <= 1)
          stop("`speed.factor` must be greater than 1")
        indxs <- indxs.init
        counter <- 0
        guesses <- bot +  ((top-bot) * (1/speed.factor)^(1:max.number.of.guesses))


        while (!checkInterval(indxs, retval=FALSE) && counter < max.number.of.guesses) {
          counter <- counter+1
          guess <- guesses[[counter]] 
          indxs <- findClosestIndx(vec, value=guess, direction="inward", beyond=indxs)
        }

        # after second pass, confirm that new indxs has tighter grip than previous index
        if (checkInterval(prev.indxs, retval=TRUE) < checkInterval(indxs, retval=TRUE))
            indxs <- prev.indxs
    }
    ci.indxs <- indxs
  }
  # ci.indxs found, one way or another
  names(ci.indxs) <- c("lower", "upper")
  
  # check if we can bring it in any. 
  ci.indx.in <- ci.indxs
  while(checkInterval(ci.indx.in, retval=FALSE)) {
    ci.indxs <- ci.indx.in
    ci.indx.in[["lower"]] <- ci.indx.in[["lower"]] + 1
  }
  ci.indx.in <- ci.indxs
  while(checkInterval(ci.indx.in, retval=FALSE)) {
    ci.indxs <- ci.indx.in
    ci.indx.in[["upper"]] <- ci.indx.in[["upper"]] - 1
  }

  return(ci.indxs)
}

findClosestIndx <- function(vec, value, direction=c("right", "left", "inward"), beyond=NA) {

  direction <- match.arg(direction)
  
  if(any(is.na(beyond))) {
    if (!missing(beyond))
      warning("NA value detected in `beyond`. Check the code that is passing a value to `findClosestIndx`")
    beyond <- length(vec) / 2
  }

  if (length(beyond)==1)
      beyond[[2]] <- beyond
  
  if (length(value)==1)
      value[[2]] <- value
  
  # initialize values
  right <- left <- (-1)

  if (direction %in% c("right", "inward") ) {
    downhill <- which(vec < value[[2]]) 
    right <- min(downhill[downhill > beyond[[2]] ])
  } 

  if (direction %in% c("left", "inward") ) {
    downhill <- which(vec < value[[1]]) 
    left <- max(downhill[downhill < beyond[[1]] ])
  } 

  ret <- c(left=left, right=right)
  ret <- ret[ret>0]
  return(ret)
}