library(pacman)

p_load(party, gbm, ranger, randomForestSRC, RLT)
source("utils.R")

index <- as.integer(Sys.getenv('SLURM_ARRAY_TASK_ID'))
run <- as.numeric(index)
set.seed(run+123456)

dat_file <- Sys.getenv('DATA_FILE')

dat <- read.table(dat_file, header=TRUE, stringsAsFactors = TRUE)

N <- NROW(dat)
if (index > N) q('n')

if (dat_file == "data/bbdata.txt") {
    role <- read.table('data/bbdsc.txt', skip = 3, stringsAsFactors = FALSE)
    xvars <- as.character(role[role$V3 %in% c("n", "c"), 2])
    yvar <- as.character(role[role$V3 == 'd', 2])
    dname <- 'baseball'
} else if (dat_file == "data/mpgdata.txt") {
    role <- read.table('data/mpgdsc.txt', skip = 3, stringsAsFactors = FALSE)
    xvars <- as.character(role[role$V3 %in% c("n", "c"), 2])
    yvar <- as.character(role[role$V3 == 'd', 2])
    dname <- 'mpg'
} else if (dat_file == "data/solder.dat") {
    role <- read.table('solder.dsc', skip = 3, stringsAsFactors = FALSE)
    colnames(dat) <- role$V2
    dat$panel <- as.factor(dat$panel)
    xvars <- as.character(role[role$V3 %in% c("n", "c"), 2])
    yvar <- as.character(role[role$V3 == 'd', 2])
    dat[,yvar] <- sqrt(dat[,yvar])
    dname <- 'solder'
} else {
    q('n')
}

fpath = paste0('all/', dname, '/res', index, '.rds')

trainx <- dat[-index,xvars]
trainy <- dat[-index,yvar]

testx <- dat[index,xvars]
testy <- dat[index,yvar]

train_df <- data.frame(trainx, y = trainy)
test_df <- data.frame(testx)

mp = length(methods) + 1

resAll <- matrix(NA, mp, 1)
rownames(resAll) <- c(names(methods), 'guide')

for (method in names(methods)) {
    method_f <- methods[[method]]
    resAll[method,1] <- method_f(trainx, trainy, test_df)
}

weights <- rep(1, N)
weights[index] <- 0

x <- dat[,xvars]
y <- dat[,yvar]

role <- ifelse(sapply(x, class) == 'factor', 'c', 'n')
train <- data.frame(x, y, w = weights)
role <- c(role, 'd', 'w')

gdir = paste0("tmp", index)
setdir = paste0("mkdir -p ", gdir)

system(setdir)

write.csv(train, file=paste0(gdir, '/tmp.txt'), row.names=FALSE, quote=TRUE)

write("tmp.txt\nNA\n2", file=paste0(gdir, '/tmp.dsc'))
write.table(cbind(1:NCOL(train), colnames(train), role), file=paste0(gdir, '/tmp.dsc'), row.names = FALSE, col.names = FALSE, quote = FALSE, append=TRUE)

system(paste0("cd ", gdir, ";../guide < ../forest.in > log.txt"), ignore.stdout = T, ignore.stderr = T)

resAll['guide',1] <- read.table(paste0(gdir, "/res.txt"), header = TRUE)[index,3]
system(paste0("rm -r ", gdir))

res <- data.frame(index = index, t(resAll))

saveRDS(res, file=fpath)
