Section 4 Machine learning model (MIEE)

library(stringr)
library(infotheo)  # For mutual information calculation
library(caret)  # For easy ensemble and other helper functions
library(randomForest)
library(dplyr)
load("./rdata/train_test_splits_PD_PDism.RData")
task = "tug"  # all,tug,cogtug
num_features = 30

# Function to perform EasyEnsemble resampling
easy_ensemble <- function(data, target, num_subsets) {
    majority_class <- names(sort(table(data[[target]]), decreasing = TRUE))[1]
    minority_class <- names(sort(table(data[[target]]), decreasing = TRUE))[2]

    majority_indices <- which(data[[target]] == majority_class)
    minority_indices <- which(data[[target]] == minority_class)

    minority_size <- length(minority_indices)

    subsets <- list()

    for (i in 1:num_subsets) {
        sampled_majority_indices <- sample(majority_indices,
            minority_size)
        subset_indices <- c(sampled_majority_indices, minority_indices)
        subsets[[i]] <- data[subset_indices, ]
    }

    return(subsets)
}

# Function to calculate mutual information and select top
# features
mi_feature_selection <- function(data, target) {
    target_var <- data[[target]]
    feature_names <- setdiff(names(data), target)

    # Calculate mutual information for each feature
    mi_scores <- sapply(feature_names, function(feature) {
        mutinformation(discretize(data[[feature]]), as.factor(target_var))
    })

    # Select top features based on MI
    mi_scores = sort(mi_scores, decreasing = TRUE)
    selected_features <- names(mi_scores)[1:num_features]

    return(selected_features)
}

# MIEE algorithm implementation
miee_algorithm <- function(data, target, num_subsets, num_trees) {
    # Step 1: Resample the dataset using EasyEnsemble
    subsets <- easy_ensemble(data, target, num_subsets)

    # Step 2: Perform feature selection using mutual
    # information on each subset
    all_selected_features <- list()
    for (i in 1:num_subsets) {
        selected_features <- mi_feature_selection(subsets[[i]],
            target)
        all_selected_features[[i]] <- selected_features
    }

    # Step 3: Train a random forest model for each subset
    models <- list()
    for (i in 1:num_subsets) {
        selected_features <- all_selected_features[[i]]
        formula <- as.formula(paste(target, "~", paste(selected_features,
            collapse = "+")))
        models[[i]] <- randomForest(formula, data = subsets[[i]],
            ntree = num_trees)
    }

    # Step 4: Ensemble predictions by averaging the
    # probability of each model
    ensemble_predict <- function(test_data) {
        predictions <- matrix(0, nrow = nrow(test_data), ncol = num_subsets)

        for (i in 1:num_subsets) {
            selected_features <- all_selected_features[[i]]
            model <- models[[i]]
            predictions[, i] <- predict(model, test_data[, selected_features],
                type = "prob")[, 2]
        }

        final_predictions <- rowMeans(predictions)
        return(ifelse(final_predictions > 0.5, "1", "0"))
    }

    return(list(models = models, predict = ensemble_predict))
}

# Build an MIEE model for each training fold
num_folds = 3
seq = 1:num_folds

for (split_num in 1:5) {
    split = all_splits[[split_num]]
    split_preds = data.frame(matrix(nrow = 0, ncol = 4))
    colnames(split_preds) = c("PD", "PDism", "PDGP", "response")

    for (i in 1:num_folds) {
        test = split[[i]]
        train_indx = seq[seq != i]
        train = data.frame(matrix(nrow = 0, ncol = length(data)))
        colnames(train) = colnames(data)
        for (indx in train_indx) {
            train = rbind(train, split[[indx]])
        }

        if (task == "tug") {
            train = train[, str_detect(colnames(train), "PDGP") |
                str_detect(colnames(train), "response") | str_detect(colnames(train),
                "\\.t4") | str_detect(colnames(train), "\\.t5") |
                str_detect(colnames(train), "\\.t4t5") | str_detect(colnames(train),
                "\\.t4t5_diff")]
        }
        predictor_vars = subset(train, select = -c(PDGP, response))
        colnames(predictor_vars) = str_replace_all(colnames(predictor_vars),
            "_turn", "_Turn")
        colnames(predictor_vars) = str_replace_all(colnames(predictor_vars),
            "_t", ".t")
        response_var = train$response

        # Convert predictors to data frame
        predictor_vars = data.frame(predictor_vars, stringsAsFactors = TRUE,
            check.names = FALSE)

        # Convert non-integer and non-numeric variables to
        # factor variables
        factor_cols = sapply(predictor_vars, function(x) class(x) !=
            "integer" & class(x) != "numeric")
        predictor_vars[, factor_cols] = data.frame(apply(predictor_vars[,
            factor_cols, drop = FALSE], 2, as.factor), stringsAsFactors = TRUE,
            check.names = FALSE)

        # Remove variables with constant values (including
        # all NAs)
        predictor_vars = predictor_vars %>%
            select(where(~n_distinct(.) > 1))

        # Impute missing values
        predictor_vars = na.roughfix(predictor_vars)

        # Remove variables having infinite values
        predictor_vars = predictor_vars[, unlist(lapply(predictor_vars,
            function(x) if (class(x) != "factor")
                is.finite(sum(x)) else TRUE))]

        # Select rows with no NAs in the response variable
        predictor_vars = predictor_vars[!is.na(response_var),
            ]
        response_var = response_var[!is.na(response_var)]

        data = cbind(predictor_vars, response_var)
        # Build MIEE model on training set
        miee_model <- miee_algorithm(data, target = "response_var",
            num_subsets = 5, num_trees = 1000)

        save(miee_model, file = paste0("./rdata/PD_PDism_miee_split",
            split_num, "_iter_", i, "_", task, "_", num_features,
            ".RData"))

        test_PDGP = test$PDGP
        # Convert predictors to data frame
        test_predictor_vars = data.frame(subset(test, select = -response),
            stringsAsFactors = TRUE, check.names = FALSE)

        # Convert non-integer and non-numeric variables to
        # factor variables
        factor_cols = sapply(test_predictor_vars, function(x) class(x) !=
            "integer" & class(x) != "numeric")
        test_predictor_vars[, factor_cols] = data.frame(apply(test_predictor_vars[,
            factor_cols, drop = FALSE], 2, as.factor), stringsAsFactors = TRUE,
            check.names = FALSE)

        # Impute missing values
        test_predictor_vars = test_predictor_vars[, !str_detect(colnames(test_predictor_vars),
            "demo")]
        test_predictor_vars = na.roughfix(test_predictor_vars)
        colnames(test_predictor_vars) = str_replace_all(colnames(test_predictor_vars),
            "_turn", "_Turn")
        colnames(test_predictor_vars) = str_replace_all(colnames(test_predictor_vars),
            "_t", ".t")
        test = cbind(test_predictor_vars, test$response)

        # Predict on test set
        predictions = miee_model$predict(test)

        preds = as.data.frame(matrix(nrow = length(predictions),
            ncol = 2))
        colnames(preds) = c("PDGP", "response")
        preds$PDGP = test_PDGP
        preds$response = predictions
        split_preds = rbind(split_preds, preds)
    }

    # Save predictions to file
    write.csv(split_preds, file = paste0("./files/split", split_num,
        "_PD_PDism_miee_", task, "_", num_features, ".csv"),
        row.names = FALSE)
}