###################################################################################################
## This script generate heterochronic datasets
## Please see associated documentation
## Authors:           A Tirado-Ramos, M Puga, C Khatchikian
## Institution:       Dartmouth College
## Publication:       BMC Reseach Notes, in review
## Date started:      Dec 2019
## Date last edited:  Dec 2022
## Version:           2.2
## Please report bugs to ckhatchikian@gmail.com
##
## Notes:   This code generates simulated heterochronous sequence datasets
##          It has 2 main sections: I) creates a phylogram/cronogram (phylogenetic tree) 
##                                  II) creates a root sequence and evolve it along the branches
##          The output consist of resulting aligments (by default in FASTA format)
##          The script is highly flexible and allow multiple user configuration
###################################################################################################


#----------------- R Packages needed -------------------------------------------------------------#
install.packages("TreeSim")       # only need to be done once
install.packages("reticulate")    # only need to be done once

library(TreeSim)
library(reticulate)


#----------------- Python Packages needed --------------------------------------------------------#
py_install(packages="pyvolve")  # only need to be done once

pyvolve<-import("pyvolve")


#----------------- Set seed (allows to replicate simulations) ------------------------------------#
set.seed(001) #if desired


#----------------- Set overall parameters for simulations ----------------------------------------#
##
## numTopo   >> numbers of phylogenetic trees (topologies)
## numRep    >> set how many replications (ie number of aligments evolved per phylogenetic tree)
### the total number of aligments in output is equal to numtopo * numRep
numRep <- 2
numTopo <- 2


##---------------- I) Tree simulation (using TreeSim) --------------------------------------------#
## Note: You can also provide the trees for sequence evolution  
## TreeSim manual can be found at https://cran.r-project.org/web/packages/TreeSim/TreeSim.pdf
##
## set the parameters (arguments) for TreeSim 
## n:       Vector of the speciation times in the incomplete phylogeny 
##          (where time is measured such that 0 is the present and increasing going into the past)
## lambda:  speciation rate 
## mu:      extincion rate 
## numTopo: number of phylogenetic trees to be simulated (defined above)


## example parameters, they can be changed (please see TreeSim manual)
n <- 10
lambda <- 2.0
mu <- 0.9

## run tree simulation, saving [numtopo] topologies simulated in t list
t <- sim.bd.taxa(n,numTopo,lambda,mu)

## assign needed attributes to trees objects
for(i in 1:numTopo){
  temp_obj <- unlist(t[i],recursive=F,use.names=T)
  attr(temp_obj,'class') <- 'phylo'
  attr(temp_obj,'order') <- 'cladewise'  
  assign(paste0("tree",i),temp_obj)
  #the line below save the trees to the working directory
  write.tree(temp_obj,paste0("tree",i,".nwk"))
  }

# To visualize individual trees, specify the trees as wanted
# example, visualize tree1 and tree2
plot(tree1)
plot(tree2)


#----------------- II) Sequence simulation (using Phyton package "pyvolve") ----------------------#
## pyvolve manual can be found at https://github.com/sjspielman/pyvolve/raw/master/user_manual/pyvolve_manual.pdf
##
## set the length (nucleotides) of the aligment
lengthNuc <- 20

## Define a nucleotide model, as a pyvolve.Model object.
nuc_model = pyvolve$Model("nucleotide")

## Assign the model to a pyvolve.Partition. The size argument was set above (lengthNuc)
partition = pyvolve$Partition(models=nuc_model,size=lengthNuc)

## Run the evolution simulation 
## NOTE: it could take some time when big aligments/trees or multiple trees/replicates are set

for(i in 1:numTopo){
  for(j in 1:numRep){
    useTree = pyvolve$read_tree(file=paste0("tree",i,".nwk"))  
    evolver = pyvolve$Evolver(partitions=partition,tree=useTree)
    seq=paste0("sequences(tree",i,"replicate",j,").fasta")
    evolver(seqfile=seq,ratefile=F,infofile=F)
  }
}