library(phytools)
library(TreeTools)

dir <- "~/Grive/Slater_Lab/Ornithoscelida_Ashley/Difficulty_estimation/"

# Find the best trees and their log-likelihood scores
get.best.ML.tree <- function(tag) {
  tmp0 <- list.files(paste0(dir, tag), pattern = "*_ML_\\d+.log", full.names = T)
  tmp1 <- Map(readLines, tmp0)
  tmp2 <- sapply(tmp1, \(x) x[grep("BEST SCORE FOUND : ", x)])
  tmp3 <- sapply(strsplit(tmp2, " : "), \(x) x[2])
  names(tmp3) <- NULL
  tmp3 <- as.numeric(tmp3)
  return(list(Best_score = max(tmp3), Tree = tmp0[which.max(tmp3)]))
}

get.best.ML.tree("BEA")
get.best.ML.tree("LEA")

# Read in the trees
bea_trees <- Map(read.tree, list.files(paste0(dir, "BEA"),
                                       pattern = "baron_ML_\\d+.treefile", full.names = T))
lea_trees <- Map(read.tree, list.files(paste0(dir, "LEA"),
                                       pattern = "langer_ML_\\d+.treefile", full.names = T))
class(bea_trees) <- "multiPhylo"
class(lea_trees) <- "multiPhylo"

# Compute the average pairwise RF distance
avg.pairwise.RF.dist <- function(treelist) {
  # Normalize by maximum possible RF distance (= 2n - 6 for two trees of n tips)
  mat <- multiRF(treelist)/(2*Ntip(treelist[[1]]) - 6)
  return(mean(mat))
}

RFall_BEA <- avg.pairwise.RF.dist(bea_trees)
RFall_LEA <- avg.pairwise.RF.dist(lea_trees)

RFall_BEA
RFall_LEA

# Get the number of unique topologies
unique.topologies <- function(treelist) {
  # Credit: Martin Smith, https://stackoverflow.com/a/65954863
  # Renumber tips in a consistent manner:
  tmp0 <- RenumberTips(treelist, treelist[[1]])
  tmp1 <- Preorder(tmp0)
  edges <- Map(\(x) x$edge, x = tmp1)
  return(length(unique(edges)))
}

Nall_BEA <- unique.topologies(bea_trees)
Nall_LEA <- unique.topologies(lea_trees)

Nall_BEA
Nall_LEA

# Get plausible trees
get.plausible.trees <- function(tag) {
  subtag <- ifelse(tag == "BEA", "baron", "langer")
  tmp0 <- readLines(paste0(dir, tag, "/", subtag, "_plausible.iqtree"))
  to_skip <- grep("Tree(\\s*)logL", tmp0)
  tmp1 <- read.table(text = tmp0, header = F, skip = to_skip + 1, nrows = 100)
  # Keep just the tree index column and the seven columns showing +/- signs (test results)
  tmp2 <- tmp1[, c(1, seq(5, ncol(tmp1), by = 2))]
  # Get the indices of those rows for which all the seven sign columns have "+"
  indices <- tmp2[apply(tmp2[, -1], 1, \(x) all(x == "+")), 1]
  # Unfortunately, R's list.files() function that produced the 'bea_trees' and 'lea_trees' lists
  # orders results differently than Unix's find utility, which was used to concatenated tree file
  # analyzed by IQ-TREE. We therefore need to reconcile their ordering:
  tmp3 <- system(paste0('cd ', dir, tag, ' && find . -name "*.treefile"'), intern = T)
  tmp4 <- gsub("\\D", "", tmp3)
  # This is the ordering produced by find/cat:
  tmp5 <- as.numeric(tmp4[tmp4 != ""])
  plaus_trees <- tmp5[indices]
  # This is the ordering produced by list.files():
  tmp6 <- as.numeric(sort(as.character(1:100)))
  new_indices <- which(tmp6 %in% plaus_trees)
  # Sanity check
  if (!identical(sort(plaus_trees), sort(tmp6[new_indices]))) {
    stop("Unable to find correct tree ordering.")
  }
  # Get the trees: nonstandard evaluation to grab 'bea_trees' or 'lea_trees' depending on the tag
  out <- eval(parse(text = paste0(tolower(tag), "_trees")))[new_indices]
  return(list(Tree_indices = plaus_trees, Trees = out))
}

bea_plausible <- get.plausible.trees("BEA")$Trees
lea_plausible <- get.plausible.trees("LEA")$Trees

RFpl_BEA <- avg.pairwise.RF.dist(bea_plausible)
RFpl_LEA <- avg.pairwise.RF.dist(lea_plausible)

RFpl_BEA
RFpl_LEA

Npl_BEA <- unique.topologies(bea_plausible)
Npl_LEA <- unique.topologies(lea_plausible)

Npl_BEA
Npl_LEA

# BEA difficulty:
(1/5) * (RFall_BEA + RFpl_BEA + (Nall_BEA/100) + (Npl_BEA/length(bea_plausible)) + (1 - (length(bea_plausible)/100)))
# LEA difficulty:
(1/5) * (RFall_LEA + RFpl_LEA + (Nall_LEA/100) + (Npl_LEA/length(lea_plausible)) + (1 - (length(lea_plausible)/100)))