set.seed(2575)
######################## Quadratic Example ########################
## Generate synthetic data with first stage budget of 25
X <- runif(25, 0,1)
Y <- X^2 + rnorm(n = length(X), sd = 0.1)

## Perform the stage one analysis using IR-wald and IR-likelihood confidence intervals 

library("twostageTE")
oneStage_IR <- stageOneAnalysis(X, Y, 0.25, type = "IR-wald", 0.99)
class(oneStage_IR)
oneStage_IR

oneStage_LR <- stageOneAnalysis(X, Y, 0.25, type = "IR-likelihood", 0.99)
oneStage_LR

## The following code chunks generate second stage data with a 
## budget of 75 and perform the second stage analysis 

## LR + LR
X2 <- seq(oneStage_LR$L1, oneStage_LR$U1, length.out = 75)
Y2 <- X2^2 + rnorm(n = length(X2), sd = 0.1)
twoStage_LR_LR <- stageTwoAnalysis(oneStage_LR, X2, Y2, 
    type = "IR-likelihood", 0.95)
twoStage_LR_LR

## IR + locLinear
X2 <- c(rep(oneStage_IR$L1,37),rep(oneStage_IR$U1,38))
Y2 <- X2^2 + rnorm(n = length(X2), sd = 0.1)
twoStage_IR_linear <- stageTwoAnalysis(oneStage_IR, explanatory = X2, 
    response = Y2, type = "locLinear", level = 0.95)
twoStage_IR_linear

## IR + IR
X2 <- seq(oneStage_IR$L1, oneStage_IR$U1, length.out = 75)
Y2 <- X2^2 + rnorm(n = length(X2), sd = 0.1)
twoStage_IR_IR <- stageTwoAnalysis(oneStage_IR, X2, Y2, type = "IR-wald", 0.95)
twoStage_IR_IR

## Combining first and second stage data
stageTwoAnalysis(oneStage_IR, X2, Y2, type = "IR-wald", 0.95,
                 combineData = TRUE)


## Summary and plot functions
summary(oneStage_IR)
summary(twoStage_IR_IR)

## Left panel of figure 3
#png("sim-quad.png")
plot(oneStage_IR) 
#dev.off()
## Right panel of figure 3
#png("sim-quad2.png")
plot(twoStage_IR_IR) 
#dev.off()

######################## Isotonic Sine Example ########################
sampleData <- function(n, lower, upper, equal = FALSE) {
  if (equal) x <- seq(lower, upper, length.out = n) 
  else x <- runif(n, lower, upper) 
  y <- (1/40) * sin(6 * pi * x) + 1/4 + x/2 + (1/4) * x^2 + rnorm(n = length(x), sd = 0.1)
  return(list(X = x, Y = y))
}

Budget <- 100
d0 <- 0.5
threshold <- (1/40) * sin(6 * pi * d0) + 1/4 + d0/2 + (1/4) * d0^2 

## Generate first stage data
n1 <- floor(Budget * 0.25)
n2 <- Budget - n1
samp <- sampleData(n1, lower = 0, upper = 1)
X <- samp$X 
Y <- samp$Y

## IR + IR
stageOne_IR <- stageOneAnalysis(X, Y, threshold, type = "IR-wald", 0.99)
samp2 <- sampleData(n2, lower = stageOne_IR$L1, upper = stageOne_IR$U1, equal = TRUE)
X2 <- samp2$X 
Y2 <- samp2$Y
twoStageIR <- stageTwoAnalysis(stageOne_IR, X2, Y2, type = "IR-wald", 0.95)

## LR + LR
stageOne_LR <- stageOneAnalysis(X, Y, threshold, type = "IR-likelihood", 0.99)
samp2 <- sampleData(n2, lower = stageOne_LR$L1, upper = stageOne_LR$U1, equal = TRUE)
X2 <- samp2$X 
Y2 <- samp2$Y
twoStageLR <- stageTwoAnalysis(stageOne_LR, X2, Y2, 
    type = "IR-likelihood", 0.95)

## IR+locLinear
X2 <- c(rep(stageOne_IR$L1,floor(n2 / 2)),rep(stageOne_IR$U1,floor(n2 / 2)))
Y2 <- (1/40) * sin(6 * pi * X2) + 1/4 + d0/2 + (1/4) * X2^2 + rnorm(n = length(X2), sd = 0.1)
twoStageLinear <- stageTwoAnalysis(stageOne_IR, X2, Y2, 
    type <- "locLinear", level <- 0.95)

## Classical single stage approach, where the entire budget is allocated in a single stage
samp <- sampleData(Budget, lower = 0, upper = 1)
X <- samp$X 
Y <- samp$Y
oneStageIR <- stageOneAnalysis(X, Y, threshold, type = "IR-wald", 0.95)
oneStageLR <- stageOneAnalysis(X, Y, threshold, type = "IR-likelihood", 0.95)

## Print results
twoStageLR
twoStageIR
twoStageLinear
oneStageLR
oneStageIR

## Left panel of figure 4
x = seq(0,1,length.out=100)
y = (1/40)*sin(6*pi*x) + 1/4 + x/2 + (1/4)*x^2
dy = (1/20)*(10*(x+1) + 3*pi*cos(6*pi*x))

#png("isosine.png")
plot(x, y, type='l', xlab="Explanatory", ylab="Response", 
ylim=c(min(y, dy), max(y, dy)), lwd=2, cex.lab=1.5)
lines(x,dy, col=2, lty=2, lwd=2)
#dev.off()

## Right panel of figure 4
#png("sim-isosine.png")
plot(twoStageLR)
#dev.off()

######################## Average Computing Times ########################
##### WARNING: THIS WILL LIKELY REQUIRE SIGNIFICANT TIME TO EXECUTE #####
library("twostageTE")

## Same isotonic sine function as before
sampleData <- function(n, lower, upper) {
  x <- runif(n, lower, upper)
  y <- (1/40) * sin(6 * pi * x) + 1/4 + x/2 + (1/4) * x^2 + rnorm(n = length(x), sd = 0.1)
  return(list(X = x, Y = y))
}
d0 <- 0.5
threshold <- (1/40) * sin(6 * pi * d0) + 1/4 + d0/2 + (1/4) * d0^2 
numReps <- 30

results <- data.frame()
counter <- 1
for (Budget in seq(100, 1000, by = 100)) { ## Overall budgets of 100, 200, ..., 1000
  cat("Starting Budget =", Budget, "\n")
  n1 <- floor(Budget * 0.25)
  n2 <- Budget - n1
  for (kk in 1:numReps) {
    cat("  kk = ", kk, "\n")
    results[counter, "Budget"] <- Budget
    samp <- sampleData(n1, lower = 0, upper = 1)
    X <- samp$X 
    Y <- samp$Y
    ## TIME OF TWO STAGE IR-WALD
    oneStage_time <- system.time(stageOneAnalysis(X, Y, threshold, type = "IR-wald", 0.99))[3]
    oneStage_IR <- stageOneAnalysis(X, Y, threshold, type = "IR-wald", 0.99)
    samp2 <- sampleData(n2, lower = oneStage_IR$L1, upper = oneStage_IR$U1)
    X2 <- samp2$X; Y2 <- samp2$Y
    results[counter, "twoStageWald"] <- oneStage_time + system.time(stageTwoAnalysis(oneStage_IR, X2, Y2, type = "IR-wald", 0.95))[3]
    ## TIME OF TWO STAGE IR-LIKELIHOOD
    oneStage_time <- system.time(stageOneAnalysis(X, Y, threshold, type = "IR-likelihood", 0.99))[3]
    oneStage_IR <- stageOneAnalysis(X, Y, threshold, type = "IR-likelihood", 0.99)
    samp2 <- sampleData(n2, lower = oneStage_IR$L1, upper = oneStage_IR$U1)
    X2 <- samp2$X ; Y2 <- samp2$Y
    results[counter, "twoStageLR"] <- oneStage_time + system.time(stageTwoAnalysis(oneStage_IR, X2, Y2, type = "IR-likelihood", 0.95))[3]		
    ## TIME OF TWO STAGE LOCLINEAR
    X2 <- c(rep(oneStage_IR$L1,37),rep(oneStage_IR$U1,38))
    Y2 <- (1/40) * sin(6 * pi * X2) + 1/4 + X2/2 + (1/4) * X2^2 + rnorm(n = length(X2), sd = 0.1)
    results[counter, "twoStageLinear"] <- oneStage_time + system.time(stageTwoAnalysis(oneStage_IR, explanatory = X2, response = Y2, type = "locLinear", level = 0.95))[3]		
    
    samp <- sampleData(Budget, lower = 0, upper = 1)
    X <- samp$X 
    Y <- samp$Y
    ## TIME OF ONE STAGE IR-WALD
    results[counter, "oneStageWald"] <- system.time(stageOneAnalysis(X, Y, threshold, type = "IR-wald", 0.95))[3]
    ## TIME OF ONE STAGE IR-LIKELIHOOD		
    results[counter, "oneStageLR"] <- system.time(stageOneAnalysis(X, Y, threshold, type = "IR-likelihood", 0.95))[3]
    
    counter <- counter + 1
  }
}
cat("\n")

library("ggplot2")
library("reshape2")
df = melt(results, id=1)

meanResults = aggregate(value ~ variable + Budget, FUN=mean, data=df)
colnames(meanResults)[1] = "Procedure"

#pdf("times-all.pdf")
ggplot(meanResults, aes(x=Budget, y=value, group=Procedure, col=Procedure, lty=Procedure, pch=Procedure)) + geom_point() + geom_line() + xlab("Budget") + ylab("Computing Time (Seconds)") + theme_bw() + theme(legend.position = "bottom")
#dev.off()