## Generate tree structure
## 1. Starts off with a bifurcate tree of
## terminal nodes n and time t 
## and the probability of speciation b = (log(n) - log(2)) / t
## The output will be a tree object:
# > names(tree) # an object of class "phylo"
# [1] "edge""Nnode"       "tip.label"   "edge.length"
## 2. Want m points in between each pair of nodes
## For each non-root node, generate m data points on the branch via
## i + runif(m) where i is the ID of a node
## and the position are upstream with respect to i.

simulate_bitree <- function(n, t, m, unify_branch = FALSE) {
  b <- (log(n) - log(2)) / t
  k <- n * 2 - 1
  tree <- phytools::pbtree(b = b, n = n, t = t, type = "discrete")
  if (unify_branch) ## Reset all branch length to 1
    tree$edge.length <- rep(1, length(tree$edge.length))
  ## Adjust tip label
  tree$tip.label <- sapply(strsplit(tree$tip.label, split='t', fixed=T), function(x) x[2])
  ## Generate positions in between nodes
  n_nodes <- tree$Nnode + length(tree$tip.label)
  positions <- vector()
  for (i in 1:n_nodes) {
    if (i != (length(tree$tip.label) + 1)) ## exclude the root here
      positions <- append(positions, i + runif(m)) 
  }
  ## Get proper node labelling
  group <- as.integer(positions)
  X_labels <- vector()
  for (i in 1:length(group)) {
    if (group[i] <= length(tree$tip.label)) {
      X_labels[i] <- tree$tip.label[group[i]]
    } else {
      X_labels[i] <- as.character(group[i])
    }
  }
  return(list(data = tree, all_nodes = positions, labels = X_labels))
}
res <- simulate_bitree(n, t, m, unify_branch)

## Generate Brownian motion on a tree and sample points from the path
## [Kushal Dey and Gao Wang]
## tree: A tree structure generated by the simulate_bitree() function res$data
## all_nodes: in-between nodes data points generated by simulate_bitree() 
## Nsim: Number of Brownian motions to generate/simulate
## diffusion: A diffusion rate parameter for the diffusion Brownian process. Default=1
## precision: The time-scale for each branch of the tree produced (if the branch length is t,
## it blows it up to t*precision and generate Brownian motion with scaled variance
## under that new scale)

## plot: if TRUE, plots the paths for each simulation,defaults to FALSE

## Samples points from Brownian path for each non-root branch
## based on tree$all_nodes() function and then records these time
## points and the corresponding value of the Tree Diffusion process
## at that point for each simulation run. This whole process is repeated
## Nsim number of times to generate Nsim Brownian processes on trees

IntermBM <- function(tree, all_nodes, Nsim, diffusion = 1, precision = 1000, plot = FALSE)
{
  Y <- list()
  ll <- floor(tree$edge.length * precision) + 1
  for (m in 1:Nsim) {
    ## create a fine grid (w.r.t precision) of discrete points to approximate continuous Brownian motion
    ## G[[i]] is finer than wnat we need (Y[[i]])
    G <- lapply(tree$edge.length,
                function(x) return(c(0, cumsum(rnorm(n = floor(x * precision), sd = sqrt(diffusion/precision))))))
    if(plot) {
      H <- nodeHeights(tree)
      plot(H[1, 1], G[[1]][1], ylim = range(G), xlim = range(H), xlab = "time", ylab = "phenotype")
      for (i in 1:length(G))
        lines(H[i, 1]:H[i, 2], G[[i]][seq(from=1, to=length(G[[i]]), length.out=length(H[i,1]:H[i,2]))])
      yy <- sapply(1:length(tree$tip.label), function(x, y) which(x == y), y = tree$edge[,2])
      yy <- sapply(yy, function(x, y) y[[x]][length(y[[x]])], y = G)
      text(x = max(H), y = yy, tree$tip.label)
    }
    ## simulate data points
    for (i in 1:nrow(tree$edge)) {
      ## find parent
      pp <- which(tree$edge[, 2] == tree$edge[i, 1])
      ## add parent's value to itself
      if (length(pp) > 0)
        ## non-root
        G[[i]] <- G[[i]] + G[[pp]][ll[pp]]
      else
        ## root
        G[[i]] <- G[[i]] + G[[1]][1]
      ## find all data points between its parent and itself
      curr_index <- which(floor(all_nodes) == tree$edge[i, 2])
      pos <- ceiling((1-all_nodes[curr_index]%%1) * length(G[[i]])) ## position: 1 minus the factional part
      ## Properly determine labels
      if (tree$edge[i, 2] <= length(tree$tip.label)) 
        label = tree$tip.label[tree$edge[i, 2]]
      else
        label = as.character(tree$edge[i, 2])
      ## Y[[label]] is a subset of G[[i]]
      if (m == 1)
        Y[[paste0('I', label)]] <- G[[i]][pos]
      else
        Y[[paste0('I', label)]] <- cbind(Y[[paste0('I', label)]], G[[i]][pos])
    }
  }
  return(Y)
}

X <- IntermBM(tree, all_nodes, Nsim, diffusion, precision)

## Merge data across nodes and perform MST on distance matrix
dat <- NULL
sample_labels <- vector()
for (i in label) {
  if (is.null(dat)) {dat <- X[[i]]
  } else {dat <- rbind(dat, X[[i]])}
  sample_labels <- append(sample_labels, rep(i, nrow(X[[i]])))
}
## use cor here for scale invariance
## projection <- ape::mst(1 - abs(cor(t(dat))))
projection <- ape::mst(dist(dat))
properties <- list(method = 'MST', samples = sample_labels, legend_text = label)

## Plot MST
## FIXME: Not sure which is the best way to visualize the MST matrix
## Here use igraph::graph.adjacency to view
source('plotconf.R')
suppressMessages(library(igraph))
fill <- rep(COLORS, ceiling(length(properties$legend_text)/length(COLORS)))[1:length(properties$legend_text)]
pdf(vizplot)
plot(graph.adjacency(X),
     edge.arrow.mode = 0,
     vertex.size = PSIZE,
     vertex.label = NA,
     vertex.color = fill[match(properties$samples, properties$legend_text)],
     vertex.frame.color = fill[match(properties$samples, properties$legend_text)],
     layout = eval(parse(text = igraph_layout))
     )
legend("right", legend = properties$legend_text, fill = fill, bty = "n", border = fill, cex = 0.75)
title(main = paste(c(properties$method, igraph_layout)))
dev.off()