#create a stochastic impuation function to easily create results over different missing rates
StochasticImpute_rep <- function(data, missing_rate)
{ 
  X.raw=data #create dataframe that is just the marker data
  X=X.raw#create a copy of the dataframe to work on
  #Set variables to be used to randomize missing values
  n=nrow(X.raw) #rows of data
  m=ncol(X.raw) #columns of data
  dp=m*n #dp gives the total # of data points - use for % calculation later
  uv=runif(dp) #uniform variable
  #We are now ready to set each missing rate
  mr1=missing_rate# missing rate (how much percentage of data is NA)
  #Converts each uniform vector into a series of TRUE/FALSE where each TRUE is a missing variable
  missing1=uv<mr1
  #Returns the true rate of missing generated to ensure function worked
  length(grep("TRUE",missing1))/length(missing1)
  #Creates a matrix of true/false values corresponding to dataframe X in layout
  index.m=matrix(missing1,n,m)#Format indicator as matrix
  #Changes "missing" (TRUE) values to read "NA" and puts them in matrix/dataframe X
  X[index.m]=NA
  n=nrow(X)
  m=ncol(X)
  fn=colSums(X,na.rm=T) # sum of genotypes for all individuals 
  fc=colSums(floor(X/3+1),na.rm=T) #count number of non missing individuals
  fa=fn/(2*fc) #Frequency of allele "2"
  for(i in 1:m){
    index.a=runif(n)<fa[i]#if above requency and na include in index 2
    index.na=is.na(X[,i])
    index.m2=index.a&index.na 
    index.m0=!index.a&index.na #if not in index 2 include in index 0
    X[index.m2,i]=2 #impute 2
    X[index.m0,i]=0 #impute 1
  }
  accuracy.r=cor(X.raw[index.m],X[index.m], use = 'complete.obs')#correlation accuracy
  index.match=X.raw==X #match index
  index.mm=index.match&index.m
  accuracy.m=length(X[index.mm])/length(X[index.m]) #proportion of match accuracy
  accsi=c(accuracy.r, accuracy.m)
  return(accsi)
}

###################################################################################################

#create a KNN impuation function to easily create results over different missing rates
knn_impute_markers <- function(data, missing_rate)
{
  X.raw<-data#create dataframe that is just the marker data
  X <- X.raw#create a copy of the dataframe to work on
  #Set variables to be used to randomize missing values
  mr=missing_rate # missing rate (how much percentage of data is NA)
  n=nrow(X)
  m=ncol(X)
  dp=m*n ##dp gives the total # of data points - use for % calculation later
  #Use KNN to impute similar to problem 3
  #Converts each uniform vector into a series of TRUE/FALSE where each TRUE is a missing variable
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)#Format indicator as matrix
  X[index.m]=NA #insert NA into index
  x=impute::impute.knn(as.matrix(t(X)))#transpose and impute
  accuracy.r=cor(X.raw[index.m],t(x$data)[index.m], use = 'complete.obs')
  index.match=X.raw==t(x$data)
  index.mm=index.match&index.m
  accuracy.m=length(X[index.mm])/length(X[index.m])
  accknn=c(accuracy.r, accuracy.m)
  return(accknn)
}


######################################################################################################

#create a KNN marker impuation function to easily create results over different missing rates
knn_geno_markers <- function(data, missing_rate)
{
  X.rawrev=t(data)#create dataframe that is just the marker data tranposed
  X.rev=X.rawrev#create a copy of the dataframe to work on
  #Sets the constants we will need (which are different than 3!)
  n=nrow(X.rawrev)
  m=ncol(X.rawrev)
  dp=m*n #dp gives the total # of data points - use for % calculation later
  #Use KNN to impute similar to problem 3
  X.rev=X.rawrev
  mr=missing_rate# missing rate (how much percentage of data is NA)
  uv=runif(dp) #Converts each uniform vector into a series of TRUE/FALSE where each TRUE is a missing variable
  missing=uv<mr
  index.m=matrix(missing,n,m)#Format indicator as matrix
  X.rev[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X.rev))) #impute transpose of transposed
  accuracy.r=cor(X.rawrev[index.m],t(x$data)[index.m], use = 'complete.obs')
  index.match=X.rawrev==t(x$data)
  index.mm=index.match&index.m
  accuracy.m=length(X.rev[index.mm])/length(X.rev[index.m])
  accknngeno=c(accuracy.r, accuracy.m)
  return(accknngeno)
}

#######################################################################################################

#create a beagle impuation function to easily create results over different missing rates
bgl_imputation <- function(data,map,missing_rate)
{
  X.raw=data#create dataframe that is just the marker data
  X=X.raw #create a copy of the dataframe to work on
  #Set variables to be used to randomize missing values
  n=nrow(X)
  m=ncol(X)
  dp=m*n #dp gives the total # of data points - use for % calculation later
  uv=runif(dp)#Converts each uniform vector into a series of TRUE/FALSE where each TRUE is a missing variable
  mr=missing_rate # missing rate (how much percentage of data is NA)
  missing=uv<mr # missing = T/F
  #Format indicator as matrix
  index.m=matrix(missing,n,m) 
  #Set missing values as NA
  X[index.m]=NA
  # Convert the genotypes to a vcf
  na_file <- "genotypes_bgl"
  vcf_na <- numeric_2_VCF(X, map)
  write_vcf(vcf_na, outfile = na_file)
  # Assign parameters
  genotype_file = "genotypes_bgl.vcf" 
  outfile = "genotypes_bgl_imp"
  
  # Define a system command
  command1_prefix <- "java -Xmx1g -jar beagle.25Nov19.28d.jar"
  command_args <- paste(" gt=", genotype_file, " out=", outfile, sep = "")
  command1 <- paste(command1_prefix, command_args)
  test <- system(command1, intern = T)
  # Run BEAGLE using the system function, this will produce a gzip .vcf file
  Xvcf=read.vcfR("genotypes_bgl_imp.vcf.gz")
  Xbgl=data.frame(Xvcf@fix,Xvcf@gt)
  genotypes_imp <- VCF_2_numeric(Xbgl)[[1]]
  # Calculate accuracy between original and imputed values
  accuracy.r=cor(X.raw[index.m],genotypes_imp[index.m], use = 'complete.obs')
  index.match=X.raw==genotypes_imp
  index.mm=index.match&index.m
  accuracy.m=length(X[index.mm])/length(X[index.m])
  accbgl=c(accuracy.r, accuracy.m)
  return(accbgl)
}