## 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) {
  tmp_n <- n
  if (n == 2) n <- 3
  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])
  if (tmp_n == 2) {
    ## A hack to simulate tree with only 2 tips
    tree$edge <- matrix(c(3,3,1,2), 2, 2)
    tree$Nnode <- 1 
    tree$tip.label <- c("1", "2")
    tree$edge.length <- tree$edge.length[3:4]
  }
  ## 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)

plot_tree <- function(tree, filename, w = 8, h = 6) {
  pdf(filename, w, h)
  ape::plot.phylo(tree, show.tip.label=F,
             show.node.label = F, edge.color="#5B9279", edge.width=4)
  ape::nodelabels(bg = "#FFFFFF", frame = 'circle')
  # labeling is tricky here ... need to specify tip.label!
  ape::tiplabels(tree$tip.label, frame = 'circle', bg = "#F0C808")
  dev.off()
}
plot_tree(tree, treeplot)

# Calculate variance/covariance matrix of data points along a tree

## 1 Create a matrix of pairwise elements of input data points
## then each entry of these matrix is covariance between a pair of data points, (i, j)
## 2. For each entry (i, j) calculation covariance:
## 2.1 when int(i) == int(j): the covariance is `adephylo::distRoot(tree, int(i)) - (max(i, j) - int(i))`
## 2.2 when i != j: `adephylo::distRoot(tree, ape::mrca(tree, full = T)[int(i), int(j)]) - foo(i - int(i), j - int(j)))`
## where foo returns: if i on j's path to MRCA then is i - int(i); if j on i's path to MRCA then is j - int(j); else zero
## It works flawlessly except for being way too slow
## \input: tree, positions
## This script works with DSC. Please load tree and positions data when you run it independently

suppressMessages(require(adephylo, quietly = T, warn.conflicts = F))
suppressMessages(require(ape, quietly = T, warn.conflicts = F))
is_subset <- function(x, y) any(apply(embed(y,length(y)-length(x)+1),2,identical,x))
get_covij <- function(tree, i, j) {
if (as.integer(i) != as.integer(j)) {
  covij <- as.vector(distRoot(tree, mrca(tree, full = T)[as.integer(i), as.integer(j)]))[1]
  path_i <- nodepath(tree, from = as.integer(i), to = length(tree$tip.label) + 1)
  path_j <- nodepath(tree, from = as.integer(j), to = length(tree$tip.label) + 1)
  ## j contains i or i contains j
  if (length(path_j) > length(path_i)) {
    covij <- covij - (i - as.integer(i)) *  is_subset(path_i, path_j)
  } else if (length(path_i) > length(path_j)) {
    covij <- covij - (j - as.integer(j)) * is_subset(path_j, path_i)
  } else {
    if (identical(path_i, path_j)) print('ERROR')
  }
} else {
  covij <- distRoot(tree, as.integer(i)) - (max(i, j) - as.integer(i))
}
return(covij)
}

CX <- matrix(NA, nrow = length(positions), ncol = length(positions))
for (ii in 1:length(positions)) {
  for (jj in 1:length(positions)) {
    if (ii >= jj) CX[ii, jj] <- get_covij(tree, positions[ii], positions[jj])
  }
}
CX[upper.tri(CX)] <- 0
CX <- CX + t(CX) - diag(diag(CX))

## Each row will be one realization from the tree,
## row count will be number of nodes on the tree
## nrow == length(X_labels)
## Each column will be on sample (replicate)
suppressMessages(require(MASS, quietly = T))
dat <- t(mvrnorm(k, rep(mu, length(X_labels)), CX * sigma * scalar))
if (type_I_noise != 0) {
  dat <- dat + matrix(rnorm(nrow(dat) * ncol(dat), sd = type_I_noise * scalar),
                      nrow = nrow(dat), ncol = ncol(dat))
}
if (type_II_noise != 0) {
  dat <- cbind(dat, t(mvrnorm(type_II_noise, rep(mu, nrow(dat)),
                              diag(nrow(dat)) * sigma * scalar)))
}
X <- list()
for (item in unique(X_labels)) {
  X[[paste0('I', item)]] <- dat[which(X_labels == item), ]
}

## 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
suppressMessages(require(MASS, quietly = T))
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)
}

AddNoise <- function(Y, noise1, noise2, sigma) {
  for (i in names(Y)) {
    if (noise1 != 0) {
      Y[[i]] <- Y[[i]] + matrix(rnorm(nrow(Y[[i]]) * ncol(Y[[i]]), sd = noise1),
                      nrow = nrow(Y[[i]]), ncol = ncol(Y[[i]]))
    }
    if (noise2 != 0) {
      Y[[i]] <- cbind(Y[[i]], t(mvrnorm(noise2, rep(0, nrow(Y[[i]])),
                              diag(nrow(Y[[i]])) * sigma)))
    }
  }
  return(Y)
}

X <- IntermBM(tree, all_nodes, Nsim, diffusion, precision)
X <- AddNoise(X, type_I_noise, type_II_noise, sqrt(diffusion/precision))

## Simulate Brownian motion along a linear tree branch
## params:
##   m: 100
##   t_scale: 1
##   Nsim: 1000
##   difussion: 100
##   precision: 1000
## return: X
## The returned data is ordered by rows (data from one end node to the other)
suppressMessages(require(MASS, quietly = T))
brown <- matrix(0, t_scale*precision, Nsim)

for (m in 1:Nsim) {
  brown[,m] <- c(0, cumsum(rnorm(n = (t_scale*precision-1), sd = sqrt(diffusion/precision))))
}

points_of_interest <- floor(seq(1, dim(brown)[1], length.out = m))
dat <- brown[points_of_interest,]
if (type_I_noise != 0) {
  dat <- dat + matrix(rnorm(nrow(dat) * ncol(dat), sd = type_I_noise),
                      nrow = nrow(dat), ncol = ncol(dat))
}
if (type_II_noise != 0) {
  dat <- cbind(dat, t(mvrnorm(type_II_noise, rep(0, nrow(dat)),
                              diag(nrow(dat)) * sqrt(diffusion/precision))))
}

X <- list(I1 = dat)

## Iteratively calling FLASH for K factors
## Because FLASH is a rank one method, each time we call FLASH, remove the variations explained and use residuals for the next round 
options(warn=-1)
source("flash.R")
flashpool <- function(data,K,
                      tol=1e-5, maxiter_r1 = 50,
                 partype = c("constant","known","Bayes_var","var_col"),
                 sigmae2_true = NA,
                 factor_value = NA,fix_factor = FALSE,
                 nonnegative = FALSE,
                 objtype = c("margin_lik","lowerbound_lik"),
                 ash_para = list(),
                 fl_list=list())
{
  res <- data
  loadings <- numeric()
  factors <- numeric()
  for(k in 1:K){
    out1 <- flash(res,
                tol=tol, 
                maxiter_r1=maxiter_r1,
                partype=partype,
                sigmae2_true = sigmae2_true,
                factor_value = factor_value,
                fix_factor = fix_factor,
                nonnegative = nonnegative,
                objtype = objtype,
                ash_para = ash_para,
                fl_list = fl_list)
    
      res <- res - out1$l%*%t(out1$f)
      loadings <- cbind(loadings, out1$l)
      factors  <- cbind(factors,  out1$f)
  }
  
  ll <- list(u=loadings,
             v=factors,
             residuals=res)
  return(ll)
}
source('utils.R')
Y <- merge_nodes(X, label)
if (is.null(K)) K <- length(Y$label) + 1
projection <- flashpool(Y$dat, K=K)
properties <- list(method = paste('Flash', 'K =', K), samples = Y$sample_labels, legend_text = Y$label)
X <- list(dat = projection$u, sample_labels = Y$sample_labels, label = Y$label)

## MST on distance matrix
source('utils.R')
Y <- merge_nodes(X, label)
## use cor here for scale invariance
## projection <- ape::mst(1 - abs(cor(t(Y$dat))))
projection <- ape::mst(dist(Y$dat))
properties <- list(method = 'MST', samples = Y$sample_labels, legend_text = Y$label)
## Give it an empty list because MST cannot be applied to itself
X <- list()

## 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()

import sys
import numpy as np
from sklearn import manifold
from sklearn.decomposition import PCA
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import minimum_spanning_tree
from plotconf import *

class DMR:
    '''
    Run various off the shelf dimensionality reduction methods
    and make figures
    '''
    def __init__(self, data, label, **kwargs):
        if label is None:
            label = list(data.keys())
        self.properties = {'range': [], 'label': label, 'color': [], 'method': None}
        self.data = None
        nb_samples = 0
        for idx, k in enumerate(label):
            if k not in data:
                continue
            self.properties['range'].append((nb_samples, nb_samples + data[k].shape[0]))
            self.properties['color'].append(COLORS[idx % len(COLORS)])
            if self.data is None:
                self.data = data[k]
            else:
                self.data = np.vstack((self.data, data[k]))
            nb_samples += data[k].shape[0]
        self.data_new = None
        self.params = kwargs

    def LLE_standard(self):
        self.properties['method'] = "LLE"
        self.data_new = manifold.LocallyLinearEmbedding(n_neighbors = self.params["n_neighbors"],
                                               n_components = self.params["n_components"],
                                               eigen_solver='auto',
                                               method = 'standard').fit_transform(self.data)

    def LLE_ltsa(self):
        self.properties['method'] = "LTSA"
        self.data_new = manifold.LocallyLinearEmbedding(n_neighbors = self.params["n_neighbors"],
                                               n_components = self.params["n_components"],
                                               eigen_solver='auto',
                                               method = 'ltsa').fit_transform(self.data)

    def LLE_hessian(self):
        self.properties['method'] = "Hessian LLE"
        self.data_new = manifold.LocallyLinearEmbedding(n_neighbors = self.params["n_neighbors"],
                                               n_components = self.params["n_components"],
                                               eigen_solver='auto',
                                               method = 'hessian').fit_transform(self.data)

    def LLE_modified(self):
        self.properties['method'] = "Modified LLE"
        self.data_new = manifold.LocallyLinearEmbedding(n_neighbors = self.params["n_neighbors"],
                                               n_components = self.params["n_components"],
                                               eigen_solver='auto',
                                               method = 'modified').fit_transform(self.data)

    def isomap(self):
        self.properties['method'] = "Isomap"
        self.data_new = manifold.Isomap(n_neighbors = self.params["n_neighbors"],
                               n_components = self.params["n_components"]).fit_transform(self.data)

    def mds(self):
        self.properties['method'] = "MDS"
        self.data_new = manifold.MDS(n_components = self.params["n_components"],
                                     max_iter=100, n_init=1).fit_transform(self.data)

    def spectral_embedding(self):
        self.properties['method'] = "Spectral Embedding"
        self.data_new = manifold.SpectralEmbedding(n_neighbors = self.params["n_neighbors"],
                                          n_components = self.params["n_components"]).fit_transform(self.data)

    def t_sne(self):
        self.properties['method'] = "t-SNE"
        self.data_new = manifold.TSNE(n_components = self.params["n_components"],
                                      init='pca', random_state=0).fit_transform(self.data)

    def pca(self):
        self.properties['method'] = "PCA"
        self.data_new = PCA(n_components = self.params["n_components"]).fit_transform(self.data)

    # def mst(self):
    #     self.properties['method'] = "MST"
    # FIXME: input is graph matrix not distance matrix
    #     self.data_new = minimum_spanning_tree(csr_matrix(1 - np.absolute(np.corrcoef(self.data)))).toarray()

dmr = DMR(data, label, n_neighbors = n_neighbors, n_components = n_comp)
try:
    getattr(dmr, sys.argv[1])()
    properties, projection = dmr.properties, dmr.data_new
except Exception as e:
    properties = projection = []
    sys.stderr.write('{}'.format(e))

# Gao Wang (c) 2016
from plotconf import *
def plot_dr_dot(data, properties, fname):
    if len(data) == 0:
        os.system('cp ~/Dropbox/.blank.pdf {}'.format(fname))
        return
    ax = plt.subplot(111)
    for x, y, z in zip(properties['range'], properties['color'], list(properties['label'])):
        plt.scatter(data[x.item(0):x.item(1), 0], data[x.item(0):x.item(1), 1],
                    color = y, label = z, s = 40)
    plt.title(properties['method'][0] + '\n')
    plt.legend(loc = "center left", bbox_to_anchor=(0.9, 0.5))
    ax.xaxis.set_major_formatter(NullFormatter())
    ax.yaxis.set_major_formatter(NullFormatter())
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    plt.axis('tight')
    plt.savefig(fname)
    plt.close()

plot_dr_dot(data, properties, vizplot)