load('/home/boyazhang/repos/unifdist/code/optim_2d_dm_6.RData') 
library(lhs)
library(laGP)
source("/home/boyazhang/repos/unifdist/code/ud_new.R")
eps <- sqrt(.Machine$double.eps)

f <- function(X)
{
  if(is.null(nrow(X))) X <- matrix(X, nrow=1)
  m <- 8.6928
  s <- 2.4269
  x1 <- 4 * X[,1] - 2
  x2 <- 4 * X[,2] - 2
  a <- 1 + (x1 + x2 + 1)^2 * (19 - 14 * x1 + 3 * x1^2 - 14 *
                                x2 + 6 * x1 * x2 + 3 * x2^2)
  b <- 30 + (2 * x1 - 3 * x2)^2 * (18 - 32 * x1 + 12 * x1^2 +
                                     48 * x2 - 36 * x1 * x2 + 27 * x2^2)
  f <- log(a * b)
  f <- (f - m)/s
  return(f)
}



EI <- function(gpi, x, fmin, pred=predGP)
{
  if(is.null(nrow(x))) x <- matrix(x, nrow=1)
  p <- pred(gpi, x, lite=TRUE)
  d <- fmin - p$mean
  if(p$s2 <= 0){
    sigma <- eps
  }else{
    sigma <- sqrt(p$s2) ## In sqrt(p$s2) : NaNs produced
  }
  #print(p)
  #print(x)
  dn <- d/sigma
  ei <- d*pnorm(dn) + sigma*dnorm(dn)
  return(ei)
}

obj.EI <- function(x, fmin, gpi) - EI(gpi, x, fmin)

EI.search <- function(X, y, gpi, multi.start=5)
{
  m <- which.min(y)
  fmin <- y[m]
  start <- matrix(X[m,], nrow=1)
  if(multi.start > 1) 
    start <- rbind(start, randomLHS(multi.start-1, ncol(X)))# start = current fmin location and multi.start-1 points from LHS design
  xnew <- matrix(NA, nrow=nrow(start), ncol=ncol(X)+1)
  for(i in 1:nrow(start)) {
    if(EI(gpi, start[i,], fmin) <= eps)
    { out <- list(value=-Inf); next }
    out <- optim(start[i,], obj.EI, method="L-BFGS-B", 
                 lower=0, upper=1, gpi=gpi, fmin=fmin)
    xnew[i,] <- c(out$par, -out$value)
  }
  if(identical(xnew, matrix(NA, nrow=nrow(start), ncol=ncol(X)+1))){
    return("didn't find xnew with positive EI")
  }
  solns <- data.frame(cbind(start, xnew))
  names(solns) <- c("s1", "s2", "x1", "x2", "val")
  solns <- solns[(solns$val > sqrt(.Machine$double.eps)),]
  return(solns)
}

optim.EI <- function(f, ninit, m, stop, theta_max, ab = NULL, design)
{
  ## set.seed(49)
  if(design == "maximin"){
    X <- maximin(ninit, m)
  } else if(design == "dist_beta"){
    X <- bd(ninit, m, T = 10^5, shape1 = 2, shape2 = 5)$X
  } else if(design == "LHS"){
    X <- randomLHS(ninit, m)
  } else if(design == "random"){
    X <- matrix(runif(ninit*m), ncol = m)
  }
  da <- darg(list(mle=TRUE, max=theta_max,  min = sqrt(.Machine$double.eps)), X) ## max = theta_max = sqrt(m)  min = sqrt(.Machine$double.eps)) 
  if(!is.null(ab)) { da$ab <- ab } ## ab = c(0,0)
  y <- f(X)
  gpi <- newGP(X, y, d=0.1, g=1e-7, dK=TRUE)
  mleGP(gpi, param="d", tmin=da$min, tmax=da$max, ab=da$ab)$msg
  maxei <- c()
  for(i in (ninit+1):stop) {
    #print(i)
    solns <- EI.search(X, y, gpi)
    while(class(solns) == "character"){
      solns <- EI.search(X, y, gpi)
      if (class(solns) != "character"){ break }
    }
    m <- which.max(solns$val)
    maxei <- c(maxei, solns$val[m])
    xnew <- as.matrix(solns[m,3:4])
    ynew <- f(xnew)
    updateGP(gpi, matrix(xnew, nrow=1), ynew)
    mle <- mleGP(gpi, param="d", tmin=da$min, tmax=da$max, ab=da$ab)
    X <- rbind(X, xnew); y <- c(y, ynew)
  }
  deleteGP(gpi)
  
  return(list(X=X, y=y, maxei=maxei))
}

## Average progress under 4 different init
## Lets repeatedly solve the problem in this way with 100 random initializations.
n <- 50
m <- 2
ninit <- 6
theta_max <- sqrt(m)*10
ab <- c(0,0)
reps <- 41
prog.ei <- matrix(NA, nrow=reps, ncol=n)
Sys.time()
for(r in 1:reps) {
  #print(r)
  os <- optim.EI(f,ninit, m, n, theta_max, ab, "dist_beta")
  prog.ei[r,1:length(os$y)] <- os$y
  for(i in 2:length(os$y)) { 
    if(is.na(prog.ei[r,i]) || prog.ei[r,i] > prog.ei[r,i-1]) 
      prog.ei[r,i] <- prog.ei[r,i-1] 
  }
}
prog.ei_1 <- prog.ei

prog.ei <- matrix(NA, nrow=reps, ncol=n)
for(r in 1:reps) {
  os <- optim.EI(f,ninit, m, n, theta_max, ab, "maximin")
  prog.ei[r,1:length(os$y)] <- os$y
  for(i in 2:length(os$y)) { 
    if(is.na(prog.ei[r,i]) || prog.ei[r,i] > prog.ei[r,i-1]) 
      prog.ei[r,i] <- prog.ei[r,i-1] 
  }
}
prog.ei_2 <- prog.ei

prog.ei <- matrix(NA, nrow=reps, ncol=n)
for(r in 1:reps) {
  os <- optim.EI(f,ninit, m, n, theta_max, ab, "LHS")
  prog.ei[r,1:length(os$y)] <- os$y
  for(i in 2:length(os$y)) { 
    if(is.na(prog.ei[r,i]) || prog.ei[r,i] > prog.ei[r,i-1]) 
      prog.ei[r,i] <- prog.ei[r,i-1] 
  }
}
prog.ei_3 <- prog.ei

prog.ei <- matrix(NA, nrow=reps, ncol=n)
for(r in 1:reps) {
  os <- optim.EI(f,ninit, m, n, theta_max, ab,  "random")
  prog.ei[r,1:length(os$y)] <- os$y
  for(i in 2:length(os$y)) { 
    if(is.na(prog.ei[r,i]) || prog.ei[r,i] > prog.ei[r,i-1]) 
      prog.ei[r,i] <- prog.ei[r,i-1] 
  }
}
prog.ei_4 <- prog.ei


Sys.time()
save.image('/home/boyazhang/repos/unifdist/code/optim_2d_dm_6.RData') 
# plot(colMeans(prog.ei_1), col=1, lwd=2, type="l")
# lines(colMeans(prog.ei_2), col=2, lwd=2, type="l")
# lines(colMeans(prog.ei_3), col=3, lwd=2, type="l")
# lines(colMeans(prog.ei_4), col=4, lwd=2, type="l")
# legend("topright", c("dist_beta", "maximin", "LHS", "random"), col=1:4, lty=1)

par(mfcol = c(1,4))
matplot(t(prog.ei_1), col="gray",lty=1, ylab="progress", xlab="its", type="l", main = "dist_beta")
matplot(t(prog.ei_2), col="gray",lty=1, ylab="progress", xlab="its", type="l", main = "maximin")
matplot(t(prog.ei_3), col="gray",lty=1, ylab="progress", xlab="its", type="l", main = "LHS")
matplot(t(prog.ei_4), col="gray",lty=1, ylab="progress", xlab="its", type="l", main = "random")

par(mfcol = c(1,1))
n <- 25
boxplot(prog.ei_1[,n], prog.ei_2[,n], prog.ei_3[,n], prog.ei_4[,n], names = c("dist_beta", "maximin", "LHS", "random"), main = paste("boxplot of fmin after", n, "searches"))

n <- 50
boxplot(prog.ei_1[,n], prog.ei_2[,n], prog.ei_3[,n], prog.ei_4[,n], names = c("dist_beta", "maximin", "LHS", "random"), main = paste("boxplot of fmin after", n, "searches"))