# code for running experiments is adapted from 
# Tagasovska, N., Chavez-Demoulin, V., & Vatter, T. (2020, November). 
# Distinguishing cause from effect using quantiles: Bivariate quantile causal discovery. 
# In International Conference on Machine Learning

source("bqcd.R")
source("baselines.R")

n_pairs <- 100


do_synthetic_benchmark <- function(method,mode,path,ext_all){
  for(ext in ext_all){
    # load ground truth for pairs
    pairs_gt <- as.numeric(read.csv(paste0(path, ext, "pairs_gt.txt"), header = FALSE)$V1)
    
    method_to_run <- switch(method,
                            "bQCCD" = QCCD,
                            "CAM" = CAMWrapper,
                            "EMD" = RKHSEWrap,
                            "LINGAM" = lingamWrapper,
                            "RESIT" = ICMLWrapper,
                            "GRAN" = GaussianityWrap,
                            "Slope" = SlopeWrap,
                            "Sloppy" = SloppyWrap,
                            "IGCI" = IGCIWrap_G,
                            "bQCD_qrf" = QCD_qrf,
                            "bQCD_qnn" = QCD_qnn)

    pair_idx <- rep("", n_pairs)
    eps <- rep(0, n_pairs)
    cds <- rep("--", n_pairs)
    time <- rep(0, n_pairs)
    eps <- rep(0, n_pairs)

    n_corr <- 0
    results <- matrix(0,nrow = n_pairs, ncol = 2)

    for(i in 1:n_pairs) {
      
      t1 <- Sys.time()
      if(mode == 1){
        pair <-  read.table(paste0(path, ext, "pair_",i,".txt"), header = FALSE,)
      } else {
        pair <-  read.table(paste0(path, ext, "pair_",i,".txt"),
                          as.is = TRUE, header = TRUE, sep = ",",
                          row.names = 1)
      }
      res <- method_to_run(pair)
      t2 <- Sys.time()
      elapsed <- as.numeric(difftime(t2, t1), units = "secs")
      pair_idx[i] <- i
      eps[i] <- res$epsilon
      cds[i] <- res$cd
      time[i] <- elapsed
      results[i,1] <- res$cd
      if( i %% 10 == 0){
        print(i)
      }
      if (!is.na(res$cd)){
        correct <- (res$cd == pairs_gt[i] )
        n_corr <- n_corr + correct
      }
    }

    correct = results[1:n_pairs,1]  == pairs_gt[1:n_pairs]
    correct[is.na(correct)] = 0
    ncorrect = sum(correct)
    corr <- rep(0, n_pairs)
    corr[pairs_gt == cds] = 1

    #resDF <- data.frame(Correct = corr, Eps = eps, Cds = cds,T = time)

    acc <- ncorrect/100
    return(acc)
  }
}

benchmark_datasets <- function(path,output_path){
  methods <- list("bQCCD","CAM","Sloppy","IGCI","RESIT")
  acc_sloppy = c()
  acc_bqccd = c()
  acc_cam = c()
  acc_igci = c()
  acc_resit = c()

  datasets <- list.dirs(path = path,full.names = FALSE)[-1]
  for(i in 1:length(datasets)){
    ext_all = c(paste(datasets[i],"/",sep=""))
    resultsAll <- lapply(methods, function (m) do_synthetic_benchmark(m,1,path,ext_all))
    resultsAll = unlist(resultsAll)
    acc_bqccd[i] = resultsAll[1]
    acc_cam[i] = resultsAll[2]
    acc_sloppy[i] = resultsAll[3]
    acc_igci[i] = resultsAll[4]
    acc_resit[i] = resultsAll[5]
  }

  resDF <- data.frame(Dataset = datasets, Sloppy = acc_sloppy,Bqccd = acc_bqccd, CAM=acc_cam, RESIT=acc_resit, IGCI=acc_igci)
  write.table(format(resDF, digits=2), file = output_path,
                  sep = "\t")
}

datasetpath = "Data/HeteroscedasticPairs/"
benchmark_datasets(datasetpath,"Results/heterosedastic_competitors.tab")

datasetpath = "Data/Linear/"
benchmark_datasets(datasetpath,"Results/linear_competitors.tab")

datasetpath = "Data/ANLSMN/"
benchmark_datasets(datasetpath,"Results/anlsmn_competitors.tab")

datasetpath = "Data/ANLS-s/"
benchmark_datasets(datasetpath,"Results/anls-s_competitors.tab")





