########## Default parameter #########
# binomial with probit link
# with increasing number of species, nlatent -> 2 - 6

if(version$minor > 5) RNGkind(sample.kind="Rounding")
library(gllvm)
load("data_sets_sparse_95.RData")
TMB::openmp(n = 6L)

result_corr_acc = result_corr_acc2= result_env = result_rmse_env =  result_time =  matrix(NA, nrow(setup),ncol = 5L)
auc = vector("list", nrow(setup))
diagnosis = vector("list", nrow(setup))

set.seed(42)

dict = as.list(2:6)
names(dict) = as.character(unique(setup$species))

accuracy = function(true, pred, t = 0.01) {
  true = true[lower.tri(true)]
  pred = pred[lower.tri(pred)]
  zero_acc = mean((abs(true) < t) == (abs(pred) < t))
  
  if(any(true>t)) {
    pos_acc = mean((true>t)[(true>t)] == (pred > t)[(true>t)])
  } else {
    pos_acc = NULL
  }
  
  if(any(true< -t)){
    neg_acc =  mean((true< -t)[(true< -t)] == (pred < -t)[(true< -t)])
  } else {
    neg_acc = NULL
  }
  return(mean(c( zero_acc, pos_acc, neg_acc)))
}

parse_cov = function(cov, t){
  cov[abs(cov) < t] = 0
  cov[cov > t] = 1
  cov[cov < -t] = -1
  return(cov)
}

accuracy2 = function(true, pred, t = 0.01){
  true = parse_cov(true, t)
  pred = parse_cov(pred, t)
  return( mean(true[lower.tri(true)] == pred[lower.tri(pred)]) )
}
counter = 1
for(i in 1:nrow(setup)) {
  sub_auc = vector("list", 5L)
  post = vector("list", 5L)
  for(j in 1:5L){
    
    tmp = data_sets[[counter]]$setup
    ### split into train and test ###
    train_X = data_sets[[counter]]$train_X
    train_Y = data_sets[[counter]]$train_Y
    test_X = data_sets[[counter]]$test_X
    test_Y = data_sets[[counter]]$test_Y
    sim = data_sets[[counter]]$sim
    counter = counter + 1L
    
    error = tryCatch({
      time = system.time({
        model = gllvm::gllvm(y = train_Y, X = data.frame(train_X),formula = ~X1+X2+X3+X4+X5, family = binomial("probit"), num.lv = dict[[as.character(tmp$species)]], seed = 42)
      })},error = function(e) e)
    if("error"  %in% class(error)) {
      rm(error)
      error = tryCatch({
        time = system.time({
          model = gllvm::gllvm(y = train_Y, X = data.frame(train_X),formula = ~X1+X2+X3+X4+X5, family = binomial("probit"), num.lv = dict[[as.character(tmp$species)]], starting.val = "zero", seed = 42)
        })},error = function(e) e)
    }
    if("error"  %in% class(error)) {
      rm(error)
      error = tryCatch({
        time = system.time({
          model = gllvm::gllvm(y = train_Y, X = data.frame(train_X),formula = ~X1+X2+X3+X4+X5, family = binomial("probit"), num.lv = dict[[as.character(tmp$species)]], starting.val = "random", seed = 42)
        })},error = function(e) e)
    }
    try({
      coefs = rbind(model$params$beta0, t(coef(model)$Xcoef))
      true_species_weights = rbind(rep(0.0, ncol(train_Y)), sim$species_weights)
      
      result_corr_acc[i,j] =  accuracy(sim$correlation, gllvm::getResidualCov(model)$cov)
      result_corr_acc2[i,j] =  accuracy2(sim$correlation, gllvm::getResidualCov(model)$cov)
      result_env[i,j] = mean(as.vector(t(coef(model)$Xcoef) > 0) == as.vector(sim$species_weights > 0))
      result_rmse_env[i,j] =  sqrt(mean((as.vector(coefs) - as.vector(true_species_weights))^2))
      result_time[i,j] = time[3]
      pred = predict.gllvm(model, newX = data.frame(test_X), type = "response")
      sub_auc[[j]] = list(pred = pred, true = test_Y)
      post[[j]] = list(correlation=gllvm::getResidualCov(model)$cov)
      rm(model)
      gc()
      .torch$cuda$empty_cache()
    },silent = TRUE)
  }
  auc[[i]] = sub_auc
  diagnosis[[i]] = post
  
  gllvm = list(
    setup = setup[i,],
    result_corr_acc = result_corr_acc,
    result_corr_acc2 = result_corr_acc2,
    result_env = result_env,
    result_rmse_env = result_rmse_env,
    result_time= result_time,
    auc = auc,
    post = diagnosis
  )
  saveRDS(gllvm, "results/6_gllvm_sparse.RDS")
}

