## apply LOCO to Diamonds data
## using random forests
## use a subsample of size 1000


rm(list=ls())

set.seed(1437)

library(randomForest)

D = read.csv("diamonds.csv",head=TRUE)
D = D[,-c(1,3,4,5)]

Y = D[,4]
D = D[,-4]
D = D[,1:3]
X = D
X = as.matrix(X)
feature_names = names(D)

options(scipen=999)

## subsample the data to speed things up
n = nrow(D)
I = sample(1:n,size=1000,replace=FALSE)
X = X[I,]
Y = Y[I]

## split into training test
n = nrow(X)
m = round(4*n/5)
I = sample(1:n,size=m,replace=FALSE)
Xtrain = X[I,]
Ytrain = Y[I]
Xtest = X[-I,]
Ytest = Y[-I]


## forest
out = randomForest(x=Xtrain,y=Ytrain,xtest=Xtest,ytest=Ytest,
                   importance=TRUE,keep.forest=TRUE)
p = predict(out,newdata=Xtest)
r = (Ytest - p)^2

## now drop each variable
P = matrix(0,length(p),3)
R = matrix(0,length(p),3)
for(j in 1:3){
     out = randomForest(x=Xtrain[,-j],y=Ytrain,xtest=Xtest[,-j],ytest=Ytest,
                        importance=TRUE,keep.forest=TRUE)
     P[,j] = predict(out,newdata=Xtest[,-j])
     R[,j] = (Ytest - P[,j])^2
     }

R = R-r
Loco = rep(0,3)
Low = rep(0,3)
Up  = rep(0,3)
n = length(p)
for(j in 1:3){
     Loco[j] = mean(R[,j])
     se = sqrt((var(R[,j])/n) + (1/n))
     Low[j] = Loco[j] - 2*se
     Up[j] = Loco[j] + 2*se
     }

print(feature_names)
print(cbind(Loco,Low,Up))

pdf("VarImpP.pdf")
varImpPlot(out)
dev.off()