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("other/", 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

resSelf <- resOther <- matrix(NA, mp, length(xvars))
colnames(resSelf) <- colnames(resOther) <- xvars
rownames(resSelf) <- rownames(resOther) <- c(names(methods), "guide")


pb <- txtProgressBar(min = 1, max = length(xvars), style = 3)
for (i in seq_along(xvars)) {
    print(xvars[i])
    xv <- xvars[i]
    
    ### ssOther
    ntrainX <- trainx[, -i]
    ntestX <- test_df[, -i]
    for (method in names(methods)) {
        method_f <- methods[[method]]
        resOther[method, xv] <- method_f(ntrainX, trainy, ntestX)
    }

    ### ssSelf
    ntrainX <- data.frame(x = trainx[, i])
    ntestX <- data.frame(x = test_df[, i])
    for (method in names(methods)) {
        method_f <- methods[[method]]
        resSelf[method, xv] <- method_f(ntrainX, trainy, ntestX)
    }

    setTxtProgressBar(pb, i)
}

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)

pb <- txtProgressBar(min = 1, max = length(xvars), style = 3)
method <- "guide"
for (i in seq_along(xvars)) {
    print(xvars[i])
    xv <- xvars[i]
    ### ssOther
    newRole <- role
    newRole[i] <- "x"

    write("tmp.txt\nNA\n2", file = paste0(gdir, "/tmp.dsc"))
    write.table(cbind(1:NCOL(train), colnames(train), newRole), 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)

    resOther[method, xv] <- read.table(paste0(gdir, "/res.txt"), header = TRUE)[index, 3]
    system(paste0("rm ", gdir, "/res.txt"))
    ### ssSelf
    newRole <- rep("x", length(xvars))
    newRole <- c(newRole, "d", "w")
    newRole[i] <- role[i]

    write("tmp.txt\nNA\n2", file = paste0(gdir, "/tmp.dsc"))
    write.table(cbind(1:NCOL(train), colnames(train), newRole), 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)

    resSelf[method, xv] <- read.table(paste0(gdir, "/res.txt"), header = TRUE)[index, 3]
    setTxtProgressBar(pb, i)
}

system(paste0("rm -r ", gdir))

other_df <- data.frame(index = index, t(resOther), feature = colnames(resOther))
self_df <- data.frame(index = index, t(resSelf), feature = colnames(resSelf))

res <- list(other = other_df, self = self_df)

saveRDS(res, file = fpath)