Section 3 Train/ Test split

library(UBL)
library(DMwR)
library(stringr)
library(dplyr)
library(randomForest)
library(caret)

set.seed(NULL)
load("./rdata/sensor_features_all_tasks.RData")

data = sensor_df
colnames(data) = str_replace_all(colnames(data), "_turn", "_Turn")
colnames(data) = str_replace_all(colnames(data), "_t", ".t")

parkinsonism_ids = c(38, 168, 194, 199, 207, 212, 214, 223, 224,
    235, 241, 259, 261, 263, 265, 268, 297, 298)
excluded_PD = c(246, 252, 293)

data = data[!data$PDGP %in% excluded_PD, ]


labels = rep(1, nrow(data))
labels[data$PDGP %in% parkinsonism_ids] = 2
data$response = as.factor(labels)

create_groups <- function(num_folds) {
    groups = list()
    PD_gp_size = ceiling(nrow(PD_data)/num_folds)
    PDism_gp_size = ceiling(nrow(PDism_data)/num_folds)

    for (i in 1:(num_folds - 1)) {
        PD_gp = sample(1:nrow(PD_data), size = PD_gp_size, replace = FALSE)
        PD_gp = PD_data[PD_gp, ]
        PD_data = PD_data[!rownames(PD_data) %in% rownames(PD_gp),
            ]

        PDism_gp = sample(1:nrow(PDism_data), size = PDism_gp_size,
            replace = FALSE)
        PDism_gp = PDism_data[PDism_gp, ]
        PDism_data = PDism_data[!rownames(PDism_data) %in% rownames(PDism_gp),
            ]

        group = rbind(PD_gp, PDism_gp)
        groups[[i]] = group
    }
    group = rbind(PD_data, PDism_data)
    groups[[num_folds]] = group
    return(groups)
}

all_splits = list()
PD_data = data[data$response == 1, ]
PDism_data = data[data$response == 2, ]
num_folds = 3

for (i in 1:5) {
    groups = create_groups(num_folds = num_folds)
    all_splits[[i]] = groups
}

data = rbind(PD_data, PDism_data)
data <- data[order(data$PDGP), ]

save(data, all_splits, file = "./rdata/train_test_splits_PD_PDism.RData")