setwd(file.path("/", "local", "groux", "scATAC-seq"))

# libraries
library(RColorBrewer)
library(motifStack)
library(TFBSTools)
library(MotifDb)

# functions
source(file.path("scripts", "functions.R"))

get_pfm_list = function(motifs, prefix_name)
{ pfm_list = list()

  for(i in 1:dim(motifs)[3])
  { pfm_list[[i]] = new("pfm",
                        mat=motifs[,,i], 
                        name=sprintf("%s %d", prefix_name, i))
  }
  return(pfm_list)
}


# number of classes searched in the data
n_classes = c(10,20,30)


for(n_class in n_classes)
{
  # get MEME motifs
  files_meme = list.files(file.path("results",
                                    "10xgenomics_PBMC_5k_peaks_meme",
                                    sprintf("%d_motifs", n_class)),
                          "mat", full.names=T)
  motifs_meme = lapply(files_meme, read.table)
  motifs_meme = lapply(motifs_meme, as.matrix)
  motifs_meme = lapply(motifs_meme, t)
  tmp = list()
  for(i in 1:length(motifs_meme))
  { rownames(motifs_meme[[i]]) = c("A", "C", "G", "T")
    tmp[[i]] = new("pfm",
                   mat=motifs_meme[[i]], 
                   name=sprintf("MEME %d", i))
  }
  motifs_meme = tmp
  rm(tmp)
  
  # get EM discovered motifs
  motifs_found = get_pfm_list(read.sequence.models(file.path("results",
                                                             "10xgenomics_PBMC_5k_peaks_classification_2",
                                                             sprintf("peaks_rmsk_sampled_sequences_1kb_%dclass_model.mat", n_class)))$models,
                              "EM")
  
  
  # colors
  red   = brewer.pal(3, "Set1")[1]
  blue  = brewer.pal(3, "Set1")[2]
  color = c(rep(blue, length(motifs_meme)),
            rep(red,  length(motifs_found)))
  # draw tree
  # X11(height=12, width=12)
  png(filename=file.path("results",
                         "10xgenomics_PBMC_5k_peaks_classification_2",
                         sprintf("%dclass_motif_tree.png", n_class)),
      units="in", res=720, width=14, height=14)
    motifStack(c(motifs_meme,
                 motifs_found),
               layout="radialPhylog", 
               circle=0.3,
               cleaves = 0.2, 
               clabel.leaves = 0.5, 
               col.bg=color,
               col.bg.alpha=0.3, 
               col.leaves=color,
               col.inner.label.circle=color, 
               inner.label.circle.width=0.05,
               col.outer.label.circle=color, 
               outer.label.circle.width=0.02, 
               circle.motif=.5,
               angle=350)

  dev.off()
}