# Author: Shirley Pu
# Code used for the paper "Examining Environmental Matrix Effects on Quantitative Non-Targeted Analysis Estimates of Per- and Polyfluoroalkyl Substances"
# Used as source file for "Pu_et_al_Env_Matrices_Supplementary_File_2.R"

### FUNCTIONS FOR INVERSE ESTIMATION WITH CALIBRATION CURVES ### 

# Function which gets the inverse estimates for a compound based on its calibration curve
# @param Cmp: String that is the name of the compound. 
# @param matrix: String that has the name of the media/matrix
# @param data_in: Input data (tibble or dataframe)
# @param leaveout: If TRUE, use leave-one-out
# @param logspace: If FALSE, convert inverse predictions from log10-space back to true concentrations
# @param more_cols: If TRUE, add more columns aside from investr() output
get_invest_summarize<-function(Cmp,matrix,data_in,leaveout=FALSE,logspace=FALSE,more_cols=TRUE){
  Cmp_data<-data_in %>% filter(Chem==Cmp,Media==matrix)
  if(leaveout){
    lms<-lapply(1:nrow(Cmp_data),function(x) lm(formula=LogArea~LogConc_uM,Cmp_data[!(1:nrow(Cmp_data) %in% x),]))
    pred_out<-sapply(1:nrow(Cmp_data),function(x) calibrate(lms[[x]],y0=Cmp_data[x,c("LogArea")] %>% as.numeric()))
  }
  else{
    one_lm<-lm(formula=LogArea~LogConc_uM,Cmp_data)
    pred_out<-sapply(Cmp_data %>% select(LogArea) %>% unlist(),function(x) calibrate(one_lm,y0=x))
  }
  pred_out<-cbind(Cmp_data,pred_out %>% t()) %>% select(-interval)
  class(pred_out$estimate)<-"numeric"
  class(pred_out$lower)<-"numeric"
  class(pred_out$upper)<-"numeric"
  if(!logspace){
    pred_out[,c("estimate","lower","upper")]<-10^(pred_out[,c("estimate","lower","upper")])
  }
  # add additional columns if more_cols==TRUE
  if(more_cols){
    if(leaveout){
      pred_out<-pred_out %>%
        mutate(Rsquared=sapply(1:nrow(pred_out),function(x) summary(lms[[x]])$r.squared))
    }
    else{
      pred_out<-pred_out %>% 
        mutate(Rsquared=summary(one_lm)$r.squared)
    }
    uM_to_ng_g<-(1e-06)*(1e9)*(1e-03)*1 # This equals 1
    pred_out<-pred_out %>% 
      # get ng/mL estimates
      mutate(estimate_ng_g=estimate*Mass_Labeled_Molar_Mass,
             lower_ng_g=lower*Mass_Labeled_Molar_Mass,
             upper_ng_g=upper*Mass_Labeled_Molar_Mass) %>%
      # Getting metrics based on ng/mL concentrations
      mutate(AQ=10^log10(estimate_ng_g/Conc)) %>% 
      mutate(AAQ=10^abs(log10(estimate_ng_g/Conc))) %>% 
      mutate(CLFR=upper_ng_g/lower_ng_g) %>%
      mutate(LCQ=Conc/lower_ng_g) %>%
      mutate(UCQ=upper_ng_g/Conc) %>%
      mutate(lower_exceeded=Conc<lower_ng_g) %>%
      mutate(upper_exceeded=Conc>upper_ng_g) %>%
      mutate(either_exceeded=Conc<lower_ng_g|Conc>upper_ng_g)
  }
  return(pred_out)
}

### FUNCTIONS FOR qNTA BOOTSTRAP APPROACHES ### 

# Function to repeat n=runs number of times to get many 2.5th, 50th, and 97.5th percentile estimates
# @param data_used: The data containing the full set of chemical RFs 
# @param training_data: The training set (the chemical RFs with the leave-one-out RFs excluded)
# @param resample_size: Set to the size of each bootstrap resample, otherwise the function defaults to resampling to the number of unique native compounds in the original full dataset
# @param perc_out: If TRUE, return only the 2.5th, 50th, and 97.5th percentile values instead of the RFs sampled and chemicals sampled
perc_fun_new<-function(data_used,training_data,resample_size=NA,perc_out=FALSE){
  # Hierarchical bootstrapping
  unique_chems<-training_data$Chem
  n_chems<-length(unique_chems)
  # Randomly sample one chemical (n*m chemicals)
  if(!is.na(resample_size)){
    chems_sampled<-sample(n_chems,size=resample_size,replace=T)
  }
  else{
    chems_sampled<-sample(n_chems,size=length(unique(data_used$Chem)),replace=T)
  }
  # Randomly sample an RF value for each chemical
  RFs_sampled<-sapply(chems_sampled,FUN=function(x) sample(training_data$RF_uM[training_data$Chem == unique_chems[x]],size=1))
  if(perc_out){
    # Take the 2.5th and 97.5th percentile RF values
    perc_2.5th<-quantile(RFs_sampled,0.025)
    perc_50th<-quantile(RFs_sampled,0.50)
    perc_97.5th<-quantile(RFs_sampled,0.975)
    # return output data frame                                      
    return(data.frame(percentile_2.5th=perc_2.5th[[1]],percentile_50th=perc_50th[[1]],percentile_97.5th=perc_97.5th[[1]]))
  }
  else{
    return(data.frame(rf_sampled=RFs_sampled,chem_sampled=unique_chems[chems_sampled]))
  }
}

# Function which does leave-one-out cross-validation with hierarchical bootstrapping
# @param data_in: Dataframe containing columns for compounds and RF values. Output from get_RF()
# @param index: The index of the chemical serving as the leave-one-out chemical
# @param loocv: If TRUE, index is used to select a chemical from the chemicals column to serve as the leave-one-out training data chemical.
# @param runs: The number of runs for the bootstrap
# @param med_out: If TRUE, only output the median values of 2.5th and 97.5th percentiles instead of the entire dataframe
# @param resample_size: Set to the size of each bootstrap resample, otherwise the function defaults to resampling to the number of unique native compounds in the original full dataset
boot_fun_new<-function(data_in,index=1,loocv=TRUE,runs,med_out=TRUE,resample_size=NA){
  # prepare the data
  cmps<-unique(data_in$Chem)
  # store the training data in a list of lists 
  training_df<-vector(mode="list",length=length(cmps))
  if(loocv){
    #test<-data_in[data_in$native==cmps[index],]
    training<-data_in[data_in$Chem != cmps[index],]
  }
  else{
    training<-data_in
  }
  #pre-allocate a dataframe
  boot_percs<-data.frame(percentile_2.5th=double(runs),percentile_50th=double(runs),percentile_97.5th=double(runs))
  # Don't need columns storing all the sampled RFs and chemicals if only outputting the median percentile values
  if(!med_out){
    boot_percs$rf_sampled=vector("list",runs)
    boot_percs$chem_sampled=vector("list",runs)
  }
  
  # do the bootstrap with n=runs number of repeats
  for(i in 1:runs){
    temp<-perc_fun_new(data_in,training,resample_size)
    boot_percs$percentile_2.5th[i]<-quantile(temp$rf_sampled,0.025)[1]
    boot_percs$percentile_50th[i]<-quantile(temp$rf_sampled,0.50)[1]
    boot_percs$percentile_97.5th[i]<-quantile(temp$rf_sampled,0.975)[1]
    if(!med_out){
      boot_percs$rf_sampled[[i]]<-temp$rf_sampled
      boot_percs$chem_sampled[[i]]<-temp$chem_sampled
    }
  }
  if(med_out){
    return(data.frame(med_2.5th_percentile=median(boot_percs$percentile_2.5th),med_50th_percentile=median(boot_percs$percentile_50th),med_97.5th_percentile=median(boot_percs$percentile_97.5th)))
  }
  else{
    return(boot_percs)
  }
}

# Function which does leave-one-out cross-validation with hierarchical bootstrapping using boot_fun_new() and returns detailed output
# @param data_in: Dataframe containing columns for native compounds and RF values. Output from get_RF()
# @param index: The index of the chemical serving as the leave-one-out chemical
# @param runs: The number of runs for the bootstrap
# @param med_out: If TRUE, only output the median values of 2.5th and 97.5th percentiles instead of the entire dataframe
# @param list_cols: If TRUE, add to the output dataframe a list-type column with all the prediction errors for a given leave-one-out chemical

loocv_new<-function(data_in,index=1,runs,med_out=TRUE,list_cols=FALSE){
  cmps<-unique(data_in$Chem)
  n<-length(cmps)
  # pre-allocate dataframe 
  loocv_out<-data.frame(min_2.5th_percentile=double(n),med_2.5th_percentile=double(n),max_2.5th_percentile=double(n),min_50th_percentile=double(n),med_50th_percentile=double(n),max_50th_percentile=double(n),min_97.5th_percentile=double(n),med_97.5th_percentile=double(n),max_97.5th_percentile=double(n),perc_out_lower=double(n),perc_in=double(n),perc_out_upper=double(n))
  # if med_out=FALSE, add the list-type columns 'rf_sampled' and 'chem_sampled' which are output in the boot_fun() dataframe
  #if(!med_out){
  #loocv_out$rf_sampled<-vector(mode="list",length=n)
  #loocv_out$chem_sampled<-vector(mode="list",length=n)
  #}
  for(i in 1:n){
    boot_out<-boot_fun_new(data_in,index=i,loocv=TRUE,runs,med_out)
    # if only medians were output
    if(med_out){
      loocv_out$med_2.5th_percentile[i]<-boot_out$med_2.5th_percentile[1]
      loocv_out$med_50th_percentile[i]<-boot_out$med_50th_percentile[1]
      loocv_out$med_97.5th_percentile[i]<-boot_out$med_97.5th_percentile[1]
    }
    # if full boot_percs dataframe was output
    else{
      loocv_out$min_2.5th_percentile[i]<-min(boot_out$percentile_2.5th)
      loocv_out$med_2.5th_percentile[i]<-median(boot_out$percentile_2.5th)
      loocv_out$max_2.5th_percentile[i]<-max(boot_out$percentile_2.5th)
      loocv_out$min_50th_percentile[i]<-min(boot_out$percentile_50th)
      loocv_out$med_50th_percentile[i]<-median(boot_out$percentile_50th)
      loocv_out$max_50th_percentile[i]<-max(boot_out$percentile_50th)
      loocv_out$min_97.5th_percentile[i]<-min(boot_out$percentile_97.5th)
      loocv_out$med_97.5th_percentile[i]<-median(boot_out$percentile_97.5th)
      loocv_out$max_97.5th_percentile[i]<-max(boot_out$percentile_97.5th)
    }
    test_rfs<-data_in$RF_uM[data_in$Chem==cmps[i]]
    loocv_out$perc_out_lower[i]<-(length(which(test_rfs<loocv_out$med_2.5th_percentile[i])))/length(test_rfs)
    loocv_out$perc_in[i]<-(length(which(test_rfs>loocv_out$med_2.5th_percentile[i] & test_rfs<loocv_out$med_97.5th_percentile[i])))/length(test_rfs)
    loocv_out$perc_out_upper[i]<-(length(which(test_rfs>loocv_out$med_97.5th_percentile[i])))/length(test_rfs)
  }
  return(loocv_out)
}

# Function which does bootstrap on the entire dataset with hierarchical bootstrapping
# @param data_in: Dataframe containing columns for native compounds and RF values. Output from get_RF()
# @param index: The index of the chemical serving as the leave-one-out chemical
# @param runs: The number of runs for the bootstrap
# @param med_out: If TRUE, only output the median values of 2.5th and 97.5th percentiles instead of the entire dataframe
# @param list_cols: If TRUE, add to the output dataframe a list-type column with all the prediction errors for a given leave-one-out chemical

full_boot<-function(data_in,index=1,runs,med_out=TRUE,list_cols=FALSE,resample_size=NA){
  # pre-allocate dataframe 
  full_boot_out<-data.frame(min_2.5th_percentile=double(1),med_2.5th_percentile=double(1),max_2.5th_percentile=double(1),min_50th_percentile=double(1),med_50th_percentile=double(1),max_50th_percentile=double(1),min_97.5th_percentile=double(1),med_97.5th_percentile=double(1),max_97.5th_percentile=double(1))
  boot_out<-boot_fun_new(data_in,index=1,loocv=FALSE,runs,med_out,resample_size)
  # if only medians were output
  if(med_out){
    full_boot_out$med_2.5th_percentile<-boot_out$med_2.5th_percentile[1]
    full_boot_out$med_50th_percentile<-boot_out$med_50th_percentile[1]
    full_boot_out$med_97.5th_percentile<-boot_out$med_97.5th_percentile[1]
  }
  # if full boot_percs dataframe was output
  else{
    full_boot_out$min_2.5th_percentile<-min(boot_out$percentile_2.5th)
    full_boot_out$med_2.5th_percentile<-median(boot_out$percentile_2.5th)
    full_boot_out$max_2.5th_percentile<-max(boot_out$percentile_2.5th)
    full_boot_out$min_50th_percentile<-min(boot_out$percentile_50th)
    full_boot_out$med_50th_percentile<-median(boot_out$percentile_50th)
    full_boot_out$max_50th_percentile<-max(boot_out$percentile_50th)
    full_boot_out$min_97.5th_percentile<-min(boot_out$percentile_97.5th)
    full_boot_out$med_97.5th_percentile<-median(boot_out$percentile_97.5th)
    full_boot_out$max_97.5th_percentile<-max(boot_out$percentile_97.5th)
  }
  return(full_boot_out)
}