#########################################
# Requirements

# Libraries
library(mvnfast)
library(gbm)
library(ICEbox)
library(iml)
library(gridExtra)
library(ggplot2)

#########################################
# 1. Simulation linear model
# y ~ x1 - x2
# n=1000
# p=2
# Corr(x1, x2) = 0

# Reproducibility
set.seed(19840)

# Data generation
X <- rmvn(n=1e3, mu=rep(0, 3), sigma=diag(3))
betaCoef <- c(1, -1, 0)
eta <- X %*% betaCoef
y <- sapply(1:length(eta), function(x) rnorm(n=1, mean=eta[x], sd=0.5) )
datFrame <- data.frame(X, y=y)

# GBM model
formulaInput <- y ~ X1 + X2 + X3
gbmFitFinal <- gbm(formula=formulaInput, distribution="gaussian",
                   data=datFrame, n.minobsinnode=1, 
                   n.trees=1000, shrinkage = 0.01,
                   interaction.depth=1,
                   n.cores=8, cv.folds=10)

# ICE
ICE_X1 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X1",
              frac_to_build = 0.25, verbose=FALSE,
              predictfcn=function(object, newdata) 
                predict(object, newdata=newdata, type="response",
                        n.trees=which.min(gbmFitFinal$cv.error)))
ICE_X2 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X2",
              frac_to_build = 0.25, verbose=FALSE,
              predictfcn=function(object, newdata) 
                predict(object, newdata=newdata, type="response",
                        n.trees=which.min(gbmFitFinal$cv.error)))
ICE_X3 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X3",
              frac_to_build = 0.25, verbose=FALSE,
              predictfcn=function(object, newdata) 
                predict(object, newdata=newdata, type="response",
                        n.trees=which.min(gbmFitFinal$cv.error)))

###########################################################
# ALE
predModel_example = Predictor$new(model = gbmFitFinal, 
                                  data=datFrame,
                                  predict.fun = function(model, newdata) {
                                    predict.gbm(object=model, newdata=newdata, 
                                                n.trees=which.min(model$cv.error))}, 
                                  type = NULL)
ALE_X1 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X1", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)
ALE_X2 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X2", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)
ALE_X3 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X3", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)





# Limits of y
yLimits <- range(ICE_X1$ice_curves, ICE_X2$ice_curves, ICE_X3$ice_curves,
                 ALE_X1$results$.ale, ALE_X2$results$.ale, ALE_X3$results$.ale)



# ALE plots
plot_ALE_X1 <- plot(ALE_X1, ylim=yLimits) + 
  theme(panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")
plot_ALE_X2 <- plot(ALE_X2, ylim=yLimits) + 
  theme(panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")
plot_ALE_X3 <- plot(ALE_X3, ylim=yLimits) + 
  theme(panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")


# Convert ICE plot to ggplot format

# X1
ggplot_base <- ggplot(data=data.frame(x=ICE_X1$gridpts,
                                      y=ICE_X1$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X1") + ylab("ICE of y") + 
  # ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))
for(i in 1:dim(ICE_X1$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X1$gridpts, 
                                                         y=ICE_X1$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X1 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X1$gridpts, 
                                                    y=ICE_X1$actual_prediction),
             size=0.5)

# X2
ggplot_base <- ggplot(data=data.frame(x=ICE_X2$gridpts,
                                      y=ICE_X2$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X2") + ylab("ICE of y") + 
  # ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))
for(i in 1:dim(ICE_X2$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X2 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X2$gridpts, 
                                                    y=ICE_X2$actual_prediction),
             size=0.5)


# X3
ggplot_base <- ggplot(data=data.frame(x=ICE_X3$gridpts,
                                      y=ICE_X3$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X3") + ylab("ICE of y") + 
  # ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))
for(i in 1:dim(ICE_X3$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X3$gridpts, 
                                                         y=ICE_X3$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X3 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X3$gridpts, y=ICE_X3$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X3$gridpts, y=ICE_X3$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X3$gridpts, 
                                                    y=ICE_X3$actual_prediction),
             size=0.5)

pdf("model1_ICE.pdf")
grid.arrange(plot_ICE_X1, 
             plot_ICE_X2, 
             plot_ICE_X3, ncol=3)
dev.off()

###################################
# 2. Linear model
# y ~ x1 - x2
# n=1000
# p=2
# Corr(x1, x2) = 0.8

# Data generation
X <- rmvn(n=1e3, mu=rep(0, 3), 
          sigma=matrix(c(1, 0.8, 0, 0.8, 1, 0, 0, 0, 1), ncol=3))
betaCoef <- c(1, -1, 0)
eta <- X %*% betaCoef
y <- sapply(1:length(eta), 
            function(x) rnorm(n=1, mean=eta[x], sd=0.5) )
datFrame <- data.frame(X, y=y)

# Check data examples
plot(datFrame[, "X1"], datFrame[, "X2"])
segments(x0=1, y0=-1, x1=1, y1=-5, lty=2)
segments(x0=1, y0=-1, x1=5, y1=-1, lty=2)
segments(x0=-1, y0=1, x1=-5, y1=1, lty=2)
segments(x0=-1, y0=1, x1=-1, y1=5, lty=2)
stopifnot(!any( (datFrame[, "X1"] > 1.25 & datFrame[, "X2"] < -1.25) |
                  (datFrame[, "X1"] < -1.25 & datFrame[, "X2"] > 1.25) ))

# GBM model
formulaInput <- y ~ X1 + X2 + X3
gbmFitFinal <- gbm(formula=formulaInput, distribution="gaussian",
                   data=datFrame, n.minobsinnode=1, 
                   n.trees=1000, shrinkage = 0.01,
                   interaction.depth=1,
                   n.cores=8, cv.folds=10)

# ICE
ICE_X1 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X1",
              frac_to_build = 0.25, verbose=TRUE,
              predictfcn=function(object, newdata) {
                ifelse( (newdata[, "X1"] > 1.25 & newdata[, "X2"] < -1.25) |
                          (newdata[, "X1"] < -1.25 & newdata[, "X2"] > 1.25) , -
                          predict(object, newdata=newdata, type="response", 
                                  n.trees=which.min(gbmFitFinal$cv.error)), 
                        predict(object, newdata=newdata, type="response",
                                n.trees=which.min(gbmFitFinal$cv.error)))
              })
ICE_X2 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X2",
              frac_to_build = 0.25, verbose=TRUE,
              predictfcn=function(object, newdata) {
                ifelse( (newdata[, "X1"] > 1.25 & newdata[, "X2"] < -1.25) |
                          (newdata[, "X1"] < -1.25 & newdata[, "X2"] > 1.25) , -
                          predict(object, newdata=newdata, type="response", 
                                  n.trees=which.min(gbmFitFinal$cv.error)), 
                        predict(object, newdata=newdata, type="response",
                                n.trees=which.min(gbmFitFinal$cv.error)))
              })
ICE_X3 <- ice(object = gbmFitFinal, X = datFrame,
              y = datFrame$y, predictor = "X3",
              frac_to_build = 0.25, verbose=TRUE,
              predictfcn=function(object, newdata) {
                ifelse( (newdata[, "X1"] > 1.25 & newdata[, "X2"] < -1.25) |
                          (newdata[, "X1"] < -1.25 & newdata[, "X2"] > 1.25) , -
                          predict(object, newdata=newdata, type="response", 
                                  n.trees=which.min(gbmFitFinal$cv.error)), 
                        predict(object, newdata=newdata, type="response",
                                n.trees=which.min(gbmFitFinal$cv.error)))
              })

# ALE
predModel_example = Predictor$new(model = gbmFitFinal, 
                                  data=datFrame,
                                  predict.fun = function(object, newdata) {
                                    ifelse( (newdata[, "X1"] > 1.25 & newdata[, "X2"] < -1.25) |
                                              (newdata[, "X1"] < -1.25 & newdata[, "X2"] > 1.25) , -
                                              predict(object, newdata=newdata, type="response", 
                                                      n.trees=which.min(gbmFitFinal$cv.error)), 
                                            predict(object, newdata=newdata, type="response",
                                                    n.trees=which.min(gbmFitFinal$cv.error)))
                                  }, 
                                  type = NULL)
ALE_X1 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X1", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)
ALE_X2 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X2", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)
ALE_X3 <- FeatureEffect$new(predictor=predModel_example, 
                            feature="X3", 
                            method = "ale", grid.size = 100,  
                            center.at = NULL)





# Limits of y
yLimits <- range(ICE_X1$ice_curves, ICE_X2$ice_curves, ICE_X3$ice_curves,
                 ALE_X1$results$.ale, ALE_X2$results$.ale, ALE_X3$results$.ale)



# ALE plots
plot_ALE_X1 <- plot(ALE_X1, ylim=yLimits) + theme(
  panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")
plot_ALE_X2 <- plot(ALE_X2, ylim=yLimits) + theme(
  panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")
plot_ALE_X3 <- plot(ALE_X3, ylim=yLimits) + theme(
  panel.background=element_rect(fill="white", colour = "black")) + ylab("ALE of y")


# Convert ICE plot to ggplot format

# X1
ggplot_base <- ggplot(data=data.frame(x=ICE_X1$gridpts,
                                      y=ICE_X1$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X1") + ylab("ICE of y") + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))
for(i in 1:dim(ICE_X1$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X1 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X1$gridpts, 
                                                    y=ICE_X1$actual_prediction),
             size=0.5)

# X2
ggplot_base <- ggplot(data=data.frame(x=ICE_X2$gridpts,
                                      y=ICE_X2$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X2") + ylab("ICE of y") + 
  # ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))

for(i in 1:dim(ICE_X2$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X2$gridpts, 
                                                         y=ICE_X2$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X2 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X2$gridpts, 
                                                    y=ICE_X2$actual_prediction),
             size=0.5)


# X3
ggplot_base <- ggplot(data=data.frame(x=ICE_X3$gridpts,
                                      y=ICE_X3$pdp),
                      mapping=aes(x=x, y=y)) + ylim(yLimits) +
  xlab("X3") + ylab("ICE of y") + 
  # ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
  theme(plot.title = element_text(hjust = 0.5),
        panel.background = element_rect(fill = "white",
                                        colour = "black"))
for(i in 1:dim(ICE_X3$ice_curves)[1]){
  ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                         data=data.frame(x=ICE_X3$gridpts, 
                                                         y=ICE_X3$ice_curves[i, ]),
                                         colour="darkgrey")
}

plot_ICE_X3 <- ggplot_base +
  geom_line(mapping=aes(x=x, y=y), 
            data=data.frame(x=ICE_X3$gridpts, y=ICE_X3$pdp), size=1,
            colour="black") +
  geom_rug(mapping=aes(x=x, y=y), 
           data=data.frame(x=ICE_X3$gridpts, y=ICE_X3$actual_prediction), 
           sides="b") +
  geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X3$gridpts, 
                                                    y=ICE_X3$actual_prediction),
             size=0.5)








# Export to PDF
pdf("model2_ICE_ALE.pdf")
grid.arrange(plot_ICE_X1, plot_ALE_X1, 
             plot_ICE_X2, plot_ALE_X2, 
             plot_ICE_X3, plot_ALE_X3, nrow=3, ncol=2)
dev.off()

########
# 3. H^2 

set.seed(19840)
h2_f_x1x2 <- function(object, datFrame){
  
  # f_x1x2
  preGrid <- data.frame(datFrame[, 1], datFrame[, 2])
  f_x1x2 <- vector("numeric", dim(preGrid)[1])
  for(j in 1:dim(preGrid)[1]){
    predGrid <- cbind(preGrid[j, ], datFrame[, 3:4], row.names = NULL)
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x1x2[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # f_x1
  preGrid <- data.frame(datFrame[, 1])
  lenPreGrid <- dim(preGrid)[1]
  f_x1 <- vector("numeric", lenPreGrid)
  for(j in 1:lenPreGrid){
    predGrid <- cbind(preGrid[j, 1], datFrame[, 2:4], row.names = NULL)
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x1[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # f_x2
  preGrid <- data.frame(datFrame[, 2])
  lenPreGrid <- dim(preGrid)[1]
  f_x2 <- vector("numeric", lenPreGrid)
  for( j in 1:lenPreGrid ){
    predGrid <- cbind(datFrame[, 1], preGrid[j, 1], datFrame[, 3:4], row.names = NULL)
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x2[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # H^2
  return( sum( (f_x1x2-f_x1-f_x2)^2 ) / sum( (f_x1x2)^2 ) )
}

h2_f_x3x4 <- function(object, datFrame){
  
  # f_x1x3
  preGrid <- data.frame(datFrame[, 3], datFrame[, 4])
  f_x3x4 <- vector("numeric", dim(preGrid)[1])
  for(j in 1:dim(preGrid)[1]){
    predGrid <- as.data.frame(cbind(datFrame[, 1:2], preGrid[j, 1:2], row.names = NULL))
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x3x4[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # f_x3
  preGrid <- data.frame(datFrame[, 3])
  lenPreGrid <- dim(preGrid)[1]
  f_x3 <- vector("numeric", lenPreGrid)
  for(j in 1:lenPreGrid){
    predGrid <- cbind(datFrame[, 1:2], preGrid[j, 1], datFrame[, 4], row.names = NULL)
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x3[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # f_x4
  preGrid <- data.frame(datFrame[, 4])
  lenPreGrid <- dim(preGrid)[1]
  f_x4 <- vector("numeric", lenPreGrid)
  for( j in 1:lenPreGrid ){
    predGrid <- cbind(datFrame[, 1:3], preGrid[j, 1], row.names = NULL)
    names(predGrid) <- c("X1", "X2", "X3", "X4")
    f_x4[j] <- mean(predict(object=object, newdata=predGrid))
  }
  
  # H^2
  return( sum( (f_x3x4-f_x3-f_x4)^2 ) / sum( (f_x3x4)^2 ) )
}

# 1. Simulate linear model
# y ~ x1 + x2 + x3 + x1 x2

X <- rmvn(n=1e2, mu=rep(0, 4), sigma=diag(4))
eta <- X[, 1] + X[, 2] + X[, 3] + X[, 1] * X[, 2]
y <- sapply(1:length(eta), function(x) rnorm(n=1, mean=eta[x], sd=0.75))

# 2. Estimate H1 model from data
lmH1 <- lm(y ~ X1 + X2 + X1:X2, data=data.frame(X, y=y))

# 3. Estimate restricted model from data
lmH0 <- lm(y ~ X1 + X2, data=data.frame(X, y=y))

# 4. Original data
H2_x1x2 <- h2_f_x1x2(lmH1, data.frame(X, y=y))
H2_x3x4 <- h2_f_x3x4(lmH1, data.frame(X, y=y))

# 5. Simulate 1e3 parametric bootstrap samples
etaBoot <- predict(lmH0, newdata=data.frame(X, y=y), type="response")
# H2boot_x1x3 <- H2boot_x1x2 <- vector("numeric", 1e3)
runFunc <- function(m, etaBoot, y){
  
  randPerm <- sample(1:length(y))
  yBoot <- etaBoot + (y[randPerm]-etaBoot[randPerm])
  
  lmH1Boot <- lm(y ~ X1 + X2 + X1:X2, 
                 data=data.frame(X, y=yBoot))
  
  # H^2
  # H2boot_x1x2[m] <- h2_f_x1x2(lmH1Boot, data.frame(X, y=yBoot))
  # H2boot_x1x3[m] <- h2_f_x1x3(lmH1Boot, data.frame(X, y=yBoot))
  H2boot_x1x2 <- h2_f_x1x2(lmH1Boot, data.frame(X, y=yBoot))
  H2boot_x3x4 <- h2_f_x3x4(lmH1Boot, data.frame(X, y=yBoot))
  return(c(H2boot_x1x2=H2boot_x1x2, H2boot_x3x4=H2boot_x3x4))
  # cat("Progress", round(m/1e3*100, 2), "%", "\n")
}

MC <- 1e3
library(parallel)
cl <- makeCluster(8)
clusterExport(cl=cl, varlist=c("y", "etaBoot", "h2_f_x1x2", "h2_f_x3x4", "X"))
simH2 <- parLapply(cl=cl, X=1:MC, fun=runFunc, y=y, etaBoot=etaBoot)
stopCluster(cl)
save(simH2, H2_x1x2, H2_x3x4, file="H2inter_H0_H1_boot")

# 6. Monte carlo p-value
load("H2inter_H0_H1_boot")
pMC_x1x2 <- (sum(sapply(simH2, function(x) x[1]) >= H2_x1x2) + 1) / 
  (MC+1)
pMC_x1x2
pMC_x3x4 <- (sum(sapply(simH2, function(x) x[2]) >= H2_x3x4) + 1) / 
  (MC+1)
pMC_x3x4

##########################
# Sim IAS

X <- rmvn(n=1e3, mu=rep(0, 3), sigma=diag(3))
A <- (X[, 1] * X[, 2])^2
B <- (X[, 1] + X[, 2] + X[, 1] * X[, 2])^2
mean(A)/mean(B) # 0.3333436

# Theoretic value 1/3

###################################
# 26.02.2020
# Simulation how PD and ICE plots change if interaction strength is increased

set.seed(19840)

# 1. Linear model
# y ~ x1 + x2
# n=1000
# p=2
# Corr(x1, x2) = 0

# Data generation
X <- rmvn(n=1e3, mu=rep(0, 2), sigma=diag(2))
betaCoef <- list(IAS0=c(1, 1, 0), 
                 IAS05=c(1, 1, sqrt(2)),
                 IAS1=c(0, 0, 1))

# Scenario 1
plot_ICE_X1 <- vector("list", 3)
plot_ICE_X2 <- vector("list", 3)
plot_ALE_X1 <- vector("list", 3)
plot_ALE_X2 <- vector("list", 3)
library(ggplot2)
library(ggplotify)
library(gbm)
library(ICEbox)
library(iml)
library(gridExtra)
for(k in 1:3){
  eta <- cbind(X, X[, 1]*X[, 2]) %*% betaCoef[[k]]
  y <- sapply(1:length(eta), function(x) rnorm(n=1, mean=eta[x], sd=0.5) )
  datFrame <- data.frame(X, y=y)
  
  # GBT model
  formulaInput <- y ~ X1 + X2
  gbmFitFinal <- gbm(formula=formulaInput, distribution="gaussian",
                     data=datFrame, n.minobsinnode=1, 
                     n.trees=1000, shrinkage = 0.01,
                     interaction.depth=2,
                     n.cores=8, cv.folds=10)
  
  # ICE
  ICE_X1 <- ice(object = gbmFitFinal, X = datFrame,
                y = datFrame$y, predictor = "X1",
                frac_to_build = 0.25, verbose=FALSE,
                predictfcn=function(object, newdata) 
                  predict(object, newdata=newdata, type="response",
                          n.trees=which.min(gbmFitFinal$cv.error)))
  ICE_X2 <- ice(object = gbmFitFinal, X = datFrame,
                y = datFrame$y, predictor = "X2",
                frac_to_build = 0.25, verbose=FALSE,
                predictfcn=function(object, newdata) 
                  predict(object, newdata=newdata, type="response",
                          n.trees=which.min(gbmFitFinal$cv.error)))
  
  # ALE
  predModel_example = Predictor$new(model = gbmFitFinal, 
                                    data=datFrame,
                                    predict.fun = function(model, newdata) {
                                      predict.gbm(object=model, newdata=newdata, 
                                                  n.trees=which.min(model$cv.error))}, 
                                    type = NULL)
  ALE_X1 <- FeatureEffect$new(predictor=predModel_example, 
                              feature="X1", 
                              method = "ale", grid.size = 100,  
                              center.at = NULL)
  ALE_X2 <- FeatureEffect$new(predictor=predModel_example, 
                              feature="X2", 
                              method = "ale", grid.size = 100,  
                              center.at = NULL)
  
  
  # Limits of y
  yLimits <- range(ICE_X1$ice_curves, ICE_X2$ice_curves,
                   ALE_X1$results$.ale, ALE_X2$results$.ale)
  
  
  # Convert ICE plot to ggplot format
  
  # X1
  ggplot_base <- ggplot(data=data.frame(x=ICE_X1$gridpts,
                                        y=ICE_X1$pdp),
                        mapping=aes(x=x, y=y)) + ylim(yLimits) +
    xlab("X1") + ylab("ICE of y") + 
    ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
    theme(plot.title = element_text(hjust = 0.5),
          panel.background = element_rect(fill = "white",
                                          colour = "black"))
  for(i in 1:dim(ICE_X1$ice_curves)[1]){
    ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                           data=data.frame(x=ICE_X1$gridpts, 
                                                           y=ICE_X1$ice_curves[i, ]),
                                           colour="darkgrey")
  }
  
  plot_ICE_X1[[k]] <- ggplot_base +
    geom_line(mapping=aes(x=x, y=y), 
              data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$pdp), size=1,
              colour="black") +
    geom_rug(mapping=aes(x=x, y=y), 
             data=data.frame(x=ICE_X1$gridpts, y=ICE_X1$actual_prediction), 
             sides="b") +
    geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X1$gridpts, 
                                                      y=ICE_X1$actual_prediction),
               size=0.5)
  
  # X2
  ggplot_base <- ggplot(data=data.frame(x=ICE_X2$gridpts,
                                        y=ICE_X2$pdp),
                        mapping=aes(x=x, y=y)) + ylim(yLimits) +
    xlab("X2") + ylab("ICE of y") + 
    ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
    theme(plot.title = element_text(hjust = 0.5),
          panel.background = element_rect(fill = "white",
                                          colour = "black"))
  for(i in 1:dim(ICE_X2$ice_curves)[1]){
    ggplot_base <- ggplot_base + geom_line(mapping=aes(x=x, y=y), 
                                           data=data.frame(x=ICE_X2$gridpts, 
                                                           y=ICE_X2$ice_curves[i, ]),
                                           colour="darkgrey")
  }
  
  plot_ICE_X2[[k]] <- ggplot_base +
    geom_line(mapping=aes(x=x, y=y), 
              data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$pdp), size=1,
              colour="black") +
    geom_rug(mapping=aes(x=x, y=y), 
             data=data.frame(x=ICE_X2$gridpts, y=ICE_X2$actual_prediction), 
             sides="b") +
    geom_point(mapping=aes(x=x, y=y), data=data.frame(x=ICE_X2$gridpts, 
                                                      y=ICE_X2$actual_prediction),
               size=0.5)
  
  
  # ALE plots
  plot_ALE_X1[[k]] <- plot(ALE_X1, ylim=yLimits) + 
    ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
    theme(plot.title = element_text(hjust = 0.5),
                                                                                                            panel.background = element_rect(fill = "white",
                                                                                                                                            colour = "black"))
  plot_ALE_X2[[k]] <- plot(ALE_X2, ylim=yLimits) + 
    ggtitle(c("IAS = 0", "IAS = 0.5", "IAS = 1")[k]) + 
    theme(plot.title = element_text(hjust = 0.5),
                                                                                                            panel.background = element_rect(fill = "white",
                                                                                                                                            colour = "black"))
}

# Export to PDF
pdf("IAS_ICEvsALE.pdf")
grid.arrange(plot_ICE_X1[[1]], plot_ICE_X2[[1]], 
             plot_ALE_X1[[1]], plot_ALE_X2[[1]],
             plot_ICE_X1[[2]], plot_ICE_X2[[2]], 
             plot_ALE_X1[[2]], plot_ALE_X2[[2]],
             plot_ICE_X1[[3]], plot_ICE_X2[[3]], 
             plot_ALE_X1[[3]], plot_ALE_X2[[3]],
             nrow=3, ncol=4)
dev.off()