####################################################################
#
# Envelope Cascade Implementation of the Bernoulli Factory
# Andrew C. Thomas, Jose H. Blanchet
# This version: June 13, 2011
#
# Includes basic execution for the form f(p)=min(multiplier*p, 1-ep).
#
####################################################################


#returns the tables necessary to get the a,b,c counts.
choose.wisely <- function (nn,kk,pp) {
  dd <- exp(sapply(kk,FUN=function(ii)
                   if ((ii>0)&(ii<nn)) sum(log(1:nn))-sum(log(1:ii))-sum(log(1:(nn-ii)))+ii*log(pp)+(nn-ii)*log(1-pp) else
                   ii*log(pp)+(nn-ii)*log(1-pp)))
  if (pp==0) dd <- c(1,rep(0,nn))
  if (pp==1) dd <- c(rep(0,nn),1)
  return(dd)
}

#Find the minimum of a quadratic-like function over an integer space. 
quadratic.integer.search <- function(fn, min.point, max.point, ...) {

  message(paste("QIS", min.point, max.point))
  retval <- NULL
  if (max.point-min.point < 10) {
    sequ <- seq(round(min.point),round(max.point))
    fn.eval <- sapply(sequ, fn, ...)
    retval <- which(fn.eval==min(fn.eval))[1]
  } else {
    ints <- ceiling((max.point-min.point)*0:4/4 + min.point)
    fn.eval <- sapply(ints, fn, ...)
    if (all(fn.eval==min(fn.eval))) retval <- ints[3] else 
    if (fn.eval[1]==min(fn.eval)) retval <- ints[1] else 
    if (fn.eval[5]==min(fn.eval)) retval <- ints[5] else 
    if (length(unique(fn.eval)) <= 3) retval <- ints[2] else 
    if (fn.eval[3]==min(fn.eval)) retval <- quadratic.integer.search(fn, ints[2], ints[4], ...) else
    if (fn.eval[2]==min(fn.eval)) retval <- quadratic.integer.search(fn, ints[1], ints[3], ...) else
    if (fn.eval[4]==min(fn.eval)) retval <- quadratic.integer.search(fn, ints[3], ints[5], ...)
  }
  return(retval)
}

#two implementations of choose(nn,kk) in the log space since we leave standard double precision very quickly.
choose.gamma <- function(nn,kk,log=FALSE) {
  r1 <- sapply(kk,FUN=function(count) lgamma(nn+1) - lgamma(nn-count+1) - lgamma(count+1))
  if (!log) r1 <- exp(r1)
  return(r1)
}
choose.ints <- function(nn,kk,log=FALSE) {
  r1 <- rep(0, length(kk))
  if (any(kk!=0)) r1[kk!=0] <- sapply(kk[kk!=0], FUN=function(count)  sum(log((nn-count+1):nn),-log(1:count)))
  r1[kk==0] <- 0
  r1[kk<0] <- -Inf
  r1[kk>nn] <- -Inf
  if (!log) r1 <- exp(r1)
  return(r1)
}
choose.sensibly <- function(nn,kk,log=FALSE) {
  r1 <- log(choose(nn,kk))
  redo <- which(r1==Inf)
  if (!log) r1 <- exp(r1)
  r1[redo] <- choose.ints(nn, kk[redo], log)
  return(r1)
}

#the main elbow function.
main.elbow <- function(pp, cc, es) pmin(cc*pp,1-es)


upper.f <- function(pp, cons) pmin(cons$cs*pp,1-cons$es) #upper envelope
bernstein.polynomial <- function (x, func=upper.f, cons) #bernstein polynomial, bounded to 1. 
    sapply(x,FUN=function(ii) min(sum(func(0:cons$n/cons$n,cons)*choose.wisely(cons$n, 0:cons$n, ii)),1))
bernstein.polynomial.unbounded <- function (x, func=upper.f, cons) 
    sapply(x,FUN=function(ii) sum(func(0:cons$n/cons$n,cons)*choose.wisely(cons$n, 0:cons$n, ii)))


#This is where the hard labor actually happens. Once these tables are defined, rberfac just makes use of them rather than recalculating each time.
berfac.preamble <- function (multiplier,
                             ep,
                             initial.n=50,
                             initial.elbow.y=(1-ep/5),
                             elbow.power=5,
                             step=500,
                             decrease.proportion=1/4,
                             total.terrace.size=8,
                             last.iteration.nn.boost=1000,
                             verbose=TRUE) {#start of preamble.

  #Diagnostic values.
  #multiplier=2; ep=0.2; initial.n=20; initial.elbow.y=0.985; step=20; total.terrace.size=3; last.iteration.nn.boost=100; verbose=TRUE; elbow.power=5; decrease.proportion=1/2;
  #prepare the counts.

  if (initial.elbow.y < 1-ep) stop("Invalid initial envelope placement.")
  
  cons <- NULL
  upper.f <- function(p, cons) pmin(cons$cs*p,1-cons$es) #upper envelope
  lower.g <- function(p, cons) upper.f(p, cons)  #pmin(multiplier*p,1-ep) #lower envelope, here the real function.
  

  lower.cons <- list(cs=multiplier, es=ep, n=initial.n)
  
  #descent curve, form-fit to small p.
  descent.x <- function(yy, pow=elbow.power, epsilon=ep, mult=multiplier) -((yy-(1-epsilon))/epsilon)^pow/multiplier + yy/multiplier
  
  f.to.min <- function(yy) abs(bernstein.polynomial(descent.x(yy), upper.f, cons)-yy)
  elbow.x <- elbow.y <- envelope.ep <- envelope.ml <- rep(NA, total.terrace.size+1)
  elbow.y[1] <- initial.elbow.y; elbow.x[1] <- descent.x(elbow.y[1])
  envelope.ep[1] <- 1-elbow.y[1]; envelope.ml[1] <- elbow.y[1]/elbow.x[1]
  
  message (paste("Initial elbow point:", elbow.x[1], elbow.y[1], ep))
  xcrit <- (1-ep)/multiplier  #critical value of x.
  ycrit <- (1-ep)

  output.sets <- NULL
  current.n <- initial.n-1
  ii <- 0
  redo.count <- 0
  iq <- 0.05
      
  repeat {
    ii <- ii+1

    cons$es <- envelope.ep[ii] <- 1-elbow.y[ii]; cons$cs <- envelope.ml[ii] <- elbow.y[ii]/elbow.x[ii];
    current.n <- current.n+1; cons$n <- current.n
    
    point.checks <- 200
    if (ii==1) {
  
      if ((bernstein.polynomial(xcrit, upper.f, cons)-ycrit<0)) {
        message (paste("For",ii,"the upper envelope cut the objective function:",
                       bernstein.polynomial(xcrit, upper.f, cons), ycrit,
                       " -- redoing previous step with more input bits."))
        redo.count <- redo.count+1
        ii <- ii - redo.count
        cons$es <- envelope.ep[ii] <- 1-elbow.y[ii]; cons$cs <- envelope.ml[ii] <- elbow.y[ii]/elbow.x[ii];
      }
    } else {
        #new criteria: pick current.n so that the gap from the curve is cut by a fixed proportion.
      current.n <- dim(output.sets[[ii-1]])[1]
      
      p1 <- bernstein.polynomial(xcrit, upper.f, cons=list(n=dim(output.sets[[ii-1]])[1]-1,
                                                   es=envelope.ep[ii-1],
                                                   cs=envelope.ml[ii-1]))
      bpn <- function(nn) bernstein.polynomial(xcrit, upper.f, cons=list(n=trunc(nn),
                                                                 es=envelope.ep[ii],
                                                                 cs=envelope.ml[ii]))
      
      test.fn <- function(chk) (chk-(decrease.proportion*ycrit+(1-decrease.proportion)*p1))^2

      first.check <- bpn(current.n)
      second.check <- bpn(current.n+1)
      if (first.check < ycrit | test.fn(second.check) > test.fn(first.check)) { ## no progress.
        ex1 <- elbow.x[ii]; ey1 <- elbow.y[ii]
        while ((bernstein.polynomial(elbow.x[ii], upper.f, cons)-(1-ep))/abs(ep-(1-elbow.y[ii-1])) > 1-decrease.proportion) {
          ex1 <- elbow.x[ii]; ey1 <- elbow.y[ii]
          
          elbow.y[ii] <- (iq*elbow.y[ii]+ycrit)/(iq+1); elbow.x[ii] <- descent.x(elbow.y[ii])
          cons$es <- 1-elbow.y[ii]; cons$cs <- elbow.y[ii]/elbow.x[ii]
        }
      
        elbow.x[ii] <- ex1;
        elbow.y[ii] <- ey1;
        cons$es <- envelope.ep[ii] <- 1-elbow.y[ii]; cons$cs <- envelope.ml[ii] <- elbow.y[ii]/elbow.x[ii]
      }

      
      temp.point.checks <- point.checks; pick <- NULL; while (length(pick)==0) {
        first.check <- bpn(current.n+(temp.point.checks-1))
        second.check <- bpn(current.n+(temp.point.checks-2))
        if (first.check < ycrit | test.fn(second.check) > test.fn(first.check)) {
          current.n <- current.n+temp.point.checks
          temp.point.checks <- temp.point.checks + point.checks
          message(paste("Minimum current.n:", current.n, ";", first.check, ycrit, p1))
        } else {
          current.n <- quadratic.integer.search(test.fn, current.n, current.n+temp.point.checks)
          pick <- current.n         
        }
      }
      cons$n <- current.n
    }



    #If this is the last round for now, push it to the limit.
    if (ii==total.terrace.size) { #milk what we can out of it!
      message(paste("Reached",ii,", the last terrace."))
      
      current.n <- current.n + last.iteration.nn.boost
      cons$n <- current.n
      ex1 <- elbow.x[ii]; ey1 <- elbow.y[ii]

      #new goal: seek the point where bf(elbow)=1-ep+1e-5).
      f.int <- function(yy) {
        xx <- descent.x(yy)
        cons <- list(es=1-yy, cs=yy/xx, n=current.n)
        return((bernstein.polynomial(xcrit, upper.f, cons)-(1-ep+1e-5))^2)
      }
      message("Optimizing final placement")
      test <- optimize(f.int, c(ey1, 1-ep))
                       #obj <- optimize(f.to.min, c(elbow.y[ii],1-ep))
      elbow.y[ii] <- test$min; elbow.x[ii] <- descent.x(elbow.y[ii])
      message(paste("Optimized final placement.", elbow.y[ii], elbow.x[ii]))

      #while (bernstein.polynomial(elbow.x[ii], upper.f, cons)-(1-ep)>0) {
      #  ex1 <- elbow.x[ii]; ey1 <- elbow.y[ii]
      #  elbow.y[ii] <- ((1-iq)*elbow.y[ii]+iq*ycrit); elbow.x[ii] <- descent.x(elbow.y[ii])
        
      #  cons$es <- 1-elbow.y[ii]; cons$cs <- elbow.y[ii]/elbow.x[ii]
      #  message(paste(bernstein.polynomial(elbow.x[ii], upper.f, cons), cons$es, cons$cs))
      #}
      #lbow.x[ii] <- ex1;
      #lbow.y[ii] <- ey1;
      cons$es <- envelope.ep[ii] <- 1-elbow.y[ii]; cons$cs <- envelope.ml[ii] <- elbow.y[ii]/elbow.x[ii]
      
    }
    lower.cons$n = current.n



    
    #Normal form.
    cs <- choose(current.n, 0:current.n)
    gs <- lower.g(0:current.n/current.n, lower.cons)
    fs <- upper.f(0:current.n/current.n, cons)

    #small fix there for machine precision.
    aa <- log(floor(cs*gs+1e-8)); bb <- log(floor(cs*(1-fs)+1e-8));
    mm <- pmax(log(cs), aa, bb);
    ccow <- mm + log(cs*exp(-mm)-exp(aa-mm)-exp(bb-mm))
#    ccow <- log(cs-aa-bb)
    #if (ii==1) print(c(aa,bb,ccow))

    bad.subs <- which(is.infinite(aa)*(aa>0) | is.infinite(bb)*(bb>0) | is.infinite(ccow)*(ccow>0) |
                      is.na(aa) | is.na(bb) | is.na(ccow))
    if (length(bad.subs)>0) { #or problems here?
      #message(paste("Invalid standard calculations in repetition",ii))
      cs <- choose.ints(current.n, (0:current.n)[bad.subs], log=TRUE)
      aa[bad.subs] <- cs+log(gs[bad.subs]); bb[bad.subs] <- cs+log(1-fs[bad.subs]);
      mm <- pmax(cs, aa[bad.subs], bb[bad.subs]); ccow[bad.subs] <- mm + log(exp(cs-mm)-exp(aa[bad.subs]-mm)-exp(bb[bad.subs]-mm))
    }
    
    output.sets[[ii]] <- cbind(aa, bb, ccow)
    colnames(output.sets[[ii]]) <- c("a","b","c")
    #writeLines (paste(ii))

    obj <- optimize(f.to.min, c(elbow.y[ii],1-ep))

    #set the next point.
    elbow.y[ii+1] <- obj$min; elbow.x[ii+1] <- descent.x(elbow.y[ii+1])

    
    #go back to the previous one. How many would be projected forward that we can remove?
    if (ii>1) {
      non <- dim(output.sets[[ii-1]])[1]-1
      stoop <- dim(output.sets[[ii]])[1]-dim(output.sets[[ii-1]])[1]
      stiff <- rbind(output.sets[[ii-1]][,1:3], array(-Inf,c(stoop,3)))
 
      newn <- non+stoop

      o2 <- output.sets[[ii]]
      for (kk in 0:newn) {
        #kk<-kk+1
        #gather all for A_{n,k}.
        precount <- choose.sensibly(stoop, 0:stoop, log=TRUE)
        a.hold <- -Inf
        #jj <- 0
        
        for (jj in 0:min(stoop,kk)) {
          #message(paste("a",stiff[kk-jj+1,1], precount[jj+1], stiff[kk-jj+1,1]+precount[jj+1], jj, kk))
          mx <- max(a.hold, stiff[kk-jj+1,1]+precount[jj+1])
          if (mx > -Inf) a.hold <- mx+log(exp(a.hold-mx)+exp(stiff[kk-jj+1,1]+precount[jj+1]-mx))
        }
        #subtract the sum total.
        #message(paste("+a",a.hold,"+",o2[kk+1,1]))
        if (a.hold > o2[kk+1,1]) {warning(paste("a.hold>o2[kk,1]", a.hold, o2[kk+1,1])); a.hold <- o2[kk+1,1]} 
        mx <- max(o2[kk+1,1], a.hold)
        if (mx > -Inf) o2[kk+1,1] <- mx + log(exp(o2[kk+1,1]-mx) - exp(a.hold-mx))
        
        b.hold <- -Inf
        for (jj in 0:min(stoop,kk)) {
          #message(paste("b",stiff[kk-jj+1,1], precount[jj+1], stiff[kk-jj+1,1]+precount[jj+1], jj, kk))
          mx <- max(b.hold, stiff[kk-jj+1,2]+precount[jj+1])
          if (mx > -Inf) b.hold <- mx+log(exp(b.hold-mx)+exp(stiff[kk-jj+1,2]+precount[jj+1]-mx))
        }
        #subtract the sum total.
        #message(paste("+b",b.hold,"+",o2[kk+1,1]))
        if (b.hold > o2[kk+1,2]) {warning(paste("b.hold>o2[kk,2]", b.hold, o2[kk+1,2])); b.hold <- o2[kk+1,2];}
        mx <- max(o2[kk+1,2], b.hold)
        if (mx > -Inf) o2[kk+1,2] <- mx + log(exp(o2[kk+1,2]-mx) - exp(b.hold-mx))
      }
      
      

      output.sets[[ii]] <- cbind(output.sets[[ii]], o2)
     
    }
    
    #output.sets[[ii]][output.sets[[ii]]<0] <- 0   #if this happens, we're in trouble with logs anyway.
    if (verbose) message(paste("Done:",ii,", bits=",nrow(output.sets[[ii]])-1))
    if (ii >= total.terrace.size) break
  }
  set.sizes <- sapply(output.sets, nrow)
  
  message(paste("Number of steps:", length(output.sets), ", sizes", paste(set.sizes,collapse=":")))
  return(list(terrace=output.sets, steps=cbind(elbow.x, elbow.y),
              multiplier=multiplier, ep=ep))
 
}




#What are the probabilities of each outcome at each n in the object, for input probability pp?
dberfac <- function(pp, terrace) {
  #terrace=stuff; pp=0.01; kk=1
  output <- sapply(1:length(terrace), function(kk) {
    prob.part <- (1:dim(terrace[[kk]])[1] - 1)*log(pp) +
      (dim(terrace[[kk]])[1]:1 - 1)*log(1-pp)
    totals <- exp(prob.part + terrace[[kk]][,1:3])
    apply(totals, 2, sum)
  })
  colnames(output) <- sapply(terrace, nrow)-1
  return(output)
}



#Given the input sequence of bits and the ``terrace'' collection of set sizes, return n.out output Bernoullis.
rberfac <- function(input,
                    n.out=1,
                    preamble=output.sets,
                    return.terrace=FALSE) {#start of rberfac.
  #preamble=stuff; input=rbinom(10000, 1, 0.01); n.out=100;return.terrace=FALSE
  terrace <- preamble$terrace
  
  p.est <- mean(input)
  stepsize <- dim(terrace[[2]])[1]-dim(terrace[[1]])[1]
  min.n <- dim(terrace[[1]])[1]-1
  curmax <- length(terrace)

  jimbo <- length(terrace)
  output <- NULL; bits.used <- NULL
  repeat {
    if (length(input)<min.n) break
    step <- 1
    tinp <- input[1:min.n]; input <- input[(min.n+1):length(input)]
    current.usage <- min.n

    kk <- sum(tinp); current.n <- length(tinp)
    drew <- terrace[[step]][kk+1,1:3]
    #message(paste(step, paste(drew, collapse=" ")))
    aset <- drew[1]; bset <- drew[2]; cset <- drew[3]
    mx <- max(aset,bset,cset)
    draw <- rbinom(1,1,exp(aset-mx)/(exp(aset-mx)+exp(bset-mx)+exp(cset-mx)))
    if (draw==0) {
      mx <- max(bset,cset)
      draw <- -rbinom(1,1,exp(cset-mx)/(exp(bset-mx)+exp(cset-mx))) #cset/(cset+bset))
    }

    if (draw>=0) {
      output <- c(output,draw)
      bits.used <- c(bits.used, current.usage)
    } else repeat {
      #if (length(input) < stepsize) break
      step <- step+1
      if (step>length(terrace)) stop (paste("rberfac went beyond the limits of the software. Input left:",length(input),
                                            "\nCurrent holding:",length(tinp),
                                            "\nCurrent outputs:",length(output),
                                            "\nOutput mean:", mean(output),
                                            "\nOnes and zeroes:", kk, current.n-kk,
                                            "\nTermination odds:", bset, cset, overdrew[2], overdrew[3]
                                            )) #terrace <- build.terrace(cons,stepsize,depth=step+3,p.est=p.est)
      stepsize <- dim(terrace[[step]])[1]-dim(terrace[[step-1]])[1]
      if (length(input) < stepsize) break
      
      tinp <- c(tinp,input[1:stepsize]); input <- input[(stepsize+1):length(input)]   #blocks of min.n in size.
      kk <- sum(tinp); current.n <- length(tinp)
      drew <- terrace[[step]][kk+1,4:6]
      overdrew <- terrace[[step]][kk+1,1:3]
      if (any(is.na(drew))) stop(paste("In stage",step,"NAs detected in the trinomial probabilities."))
      #message(paste(step, paste(drew, collapse=" ")))

      aset <- drew[1]; bset <- drew[2]; cset <- drew[3]
      
      mx <- max(aset,bset,cset)
      draw <- rbinom(1,1,exp(aset-mx)/(exp(aset-mx)+exp(bset-mx)+exp(cset-mx)))
      if (draw==0) {
        mx <- max(bset,cset)
        draw <- -rbinom(1,1,exp(cset-mx)/(exp(bset-mx)+exp(cset-mx))) #cset/(cset+bset))
      }
      #draw <- rbinom(1,1,aset/(aset+bset+cset)); if (draw==0) draw <- -rbinom(1,1,cset/(cset+bset))

      if (draw>=0) {
        output <- c(output,draw);
        bits.used <- c(bits.used, current.n)
        break
      }
      if (length(input)<min.n) break
    }

    if (length(input)<min.n) break
    if (length(output)>=n.out) break
  }

  #print(length(output))
  if (length(output)<n.out) warning("Ran out of input bits before requested number of outputs was reached.")
  
  out.list <- list(output=output, bits.used=bits.used, unused.input=input)
  if (return.terrace) out.list$terrace <- terrace
  return(out.list)
}



#Nice output of the cascade curves.
plot.berfac.curves <- function(output,  pointcount=500, quick=TRUE, descent.curve=TRUE, elbow.power=20, ...) {
  #output=stuff.5; pointcount=50
  multiplier=output$mult; ep=output$ep; 
  descent.x <- function(yy, pow=elbow.power, epsilon=ep, mult=multiplier) -((yy-(1-epsilon))/epsilon)^pow/multiplier + yy/multiplier

  #main.
  plot(c(0,(1-ep/multiplier)*1.1), c(1-2*ep,1), ty="n", main="Bernoulli Factory Output", ...)
  points(output$steps)
  points((1-ep)/multiplier, 1-ep, col=4, pch=19)

  yy <- seq(1, 1-ep, length=100)
  if (descent.curve) lines(descent.x(yy), yy, lty=3, lwd=2)
    
  xx <- 0:pointcount/pointcount
  lines (xx, upper.f(xx, list(cs=multiplier, es=ep)), col=3, lwd=2, lty=2)
  
  #Censor the output at the top, for easy visualization; for most cases this is only an issue for low n at the top edge.
  out.yy <- sapply(xx, function(kk) pmin(1-dberfac(kk, output$terrace)[2,], output$steps[,2]))
  out.yy[,1] <- 0; out.yy[,pointcount+1] <- out.yy[,pointcount]
  
  for (kk in 1:length(output$terrace)) {
    nn <- dim(output$terrace[[kk]])[1]-1; print(nn)
    lines (c(0, output$steps[kk,1], 1), c(0, output$steps[kk,2], output$steps[kk,2]))
    lines (xx, out.yy[kk,], col=kk+1)
  }

}


#Single visual output of the factory.
plot.b1 <- function(nn, multiplier, ep) {
  xx <- 0:100/100
  lines (xx, upper.f(xx, list(cs=multiplier, es=ep)), col=3, lwd=2)
  lines (xx, bernstein.polynomial (xx, upper.f,
                                   list(es=ep,
                                        cs=multiplier,
                                        n=nn)),
         col=4)

}


#Given an output terrace, simulate 10000 trials for the factory.
ex.val <- function(output, pp=0.01, trials=10000) {
  o1 <- dberfac(pp, output$terrace)
  nn <- as.numeric(colnames(o1))
  nn <- c(nn, 1000*nn[length(nn)]) #for now, make the number of draws for ``failure'' huge.
  surv <- o1[3,]
  probs <- c(1, surv)-c(surv,0)
  draws <- sample(nn, trials, prob=probs, replace=TRUE)
  if (any(draws==nn[length(nn)])) warning ("The empirical expected value includes rough estimates from empirically failed trials.")
  return(c(mean(draws), sd(draws)))
}


demos <- function() {

#  q("no")
#  source("rberfac-public.R")
  stuff <- berfac.preamble(multiplier=2, ep=0.2, initial.n=20,
                           step=20, total=3, decrease.proportion=1/2,
                           last=1000,
                           initial.elbow=0.985, elbow.power=5)
  
  dberfac(0.01, stuff$terrace)
  rberfac(input=rbinom(100000, 1, 0.01), 10000, stuff)
  save(stuff, file="bf2.RData")
  load("bf2.RData")
  ex.val(stuff)
  
  pdf ("cascade2.pdf", width=6, height=6)
  plot.berfac.curves(stuff, quick=FALSE, xlab="p", ylab="f(p)", elbow.power=5)
  dev.off()



}