# This script will perform exploratory analysis and stochastic imputation using the GAPIT demo data

# Load dependencies
library(ggplot2)
library(cvTools) # This package will be used for randomized k-fold cross-validation of imputation accuracy
source('IMPUTATION_functions.R')

# Import genotypes
gt_scan <- data.frame(read.table("mdp_numeric.txt", header = T, stringsAsFactors = F, sep = "\t", nrows = 1))
classes <- sapply(gt_scan, class)
genotypes <- data.frame(read.table("mdp_numeric.txt", header = T, row.names = 1, colClasses = classes, stringsAsFactors = F, sep = "\t"))

marker_map <- read.table("mdp_numeric.map", header = T, stringsAsFactors = F, sep = "\t")

# Check if there are any missing values
unique(unlist(genotypes))

# Check marker minor allele frequencies
minor_allele_freq <- apply(genotypes, 2, function(x)
  {
  allele_freq1 <- (sum(x == 0)*2 + sum(x == 1)) / (sum(!is.na(x)) * 2)
  allele_freq2 <- (sum(x == 2)*2 + sum(x == 1)) / (sum(!is.na(x)) * 2)
  
  return(min(allele_freq1, allele_freq2))
})

# Plot the minor allele frequency
qplot(minor_allele_freq, bins = 50) +
  labs(x = 'Minor Allele Frequency', y = "Number of Genes", title = 'Histogram of minor allele frequency')

# Filter out markers with maf <= 0.05
genotypes_maf_filt <- genotypes[,minor_allele_freq > 0.05]

# Split the genotype file into 100 folds (each fold will be ~1% of the data)
folds <- cvFolds(length(unlist(genotypes)), K = 100)
genotypes_na_folds <- get_gen_cv_sets(folds, genotypes, fold_ct = 10)

# Testing just a single fold
genotypes_na <- genotypes_na_folds[[1]][[1]] # Extract just a single fold for easy testing
genotypes_imp <- impute_mean(genotypes_na) # Impute the values
accuracy <- calc_impute_acc(genotypes, genotypes_na, genotypes_imp) # Calculate the accuracy

# Testing across all folds
accuracy_cv <- test_mean_impute_cv(genotypes, genotypes_na_folds)

# Plot these accuracies
qplot(x = factor(1), y = accuracy_cv, geom = 'boxplot') +
  labs(title = 'Accuracy of mean imputation', x = 'Missing percentage', y = 'Accuracy')
