################################################################################
#
# RevBayes Example: Bayesian inference of phylogeny using a GTR+Gamma+Inv
#                   substitution model for a single, uniform partition model.
#
# authors: Michael Landis, Sebastian Hoehna, Tracy A. Heath and Brian R. Moore
#
################################################################################

###### This just defines a single model for all sites #######

### Read in sequence data for both genes
data_cox2 = readDiscreteCharacterData("data/primates_and_galeopterus_cox2.nex")
data_cytb = readDiscreteCharacterData("data/primates_and_galeopterus_cytb.nex")

### Create concatenated data matrix
data = concatenate( data_cox2, data_cytb )

# Get some useful variables from the data. We need these later on.
n_species <- data.ntaxa()
n_branches <- 2 * n_species - 3
taxa <- data.taxa()

mvi = 0 
mni = 0


######################
# Substitution Model #
######################

#### specify the GTR+G+I substitution model applied uniformly to all sites ###
er_prior <- v(1,1,1,1,1,1)
er ~ dnDirichlet(er_prior)
moves[++mvi] = mvBetaSimplex(er, alpha=10.0, tune=true, weight=3.0)
moves[++mvi] = mvDirichletSimplex(er, alpha=10.0, tune=true, weight=1.0)

pi_prior <- v(1,1,1,1) 
pi ~ dnDirichlet(pi_prior)
moves[++mvi] = mvBetaSimplex(pi, alpha=10.0, tune=true, weight=2.0)
moves[++mvi] = mvDirichletSimplex(pi, alpha=10.0, tune=true, weight=1.0)


#### create a deterministic variable for the rate matrix ####
Q := fnGTR(er,pi) 


#############################
# Among Site Rate Variation #
#############################

alpha_prior_mean <- ln(2.0)
alpha_prior_sd <- 0.587405
alpha ~ dnLognormal( alpha_prior_mean, alpha_prior_sd )
gamma_rates := fnDiscretizeGamma( alpha, alpha, 4, false )

# add moves for the stationary frequencies, exchangeability rates and the shape parameter
moves[++mvi] = mvScale(alpha,weight=2)


# the probability of a site being invariable
pinvar ~ dnBeta(1,1)
moves[++mvi] = mvScale(pinvar, lambda=0.1, tune=true, weight=2.0)
moves[++mvi] = mvSlide(pinvar, delta=10.0, tune=true, weight=2.0)



##############
# Tree model #
##############

out_group = clade("Galeopterus_variegatus")
# Prior distribution on the tree topology
topology ~ dnUniformTopology(taxa, outgroup=out_group)
moves[++mvi] = mvNNI(topology, weight=5.0)
moves[++mvi] = mvSPR(topology, weight=1.0)

# Branch length prior
for (i in 1:n_branches) {
    bl[i] ~ dnExponential(10.0)
    moves[++mvi] = mvScale(bl[i])
}

TL := sum(bl)

psi := treeAssembly(topology, bl)




###################
# PhyloCTMC Model #
###################

# the sequence evolution model
phyloSeq ~ dnPhyloCTMC(tree=psi, Q=Q, siteRates=gamma_rates, pInv=pinvar, type="DNA")

# attach the data
phyloSeq.clamp(data)


############
# Analysis #
############

mymodel = model(psi)

# add monitors
monitors[++mni] = mnScreen(alpha, pinvar, TL, printgen=1000)
monitors[++mni] = mnFile(psi, filename="output/PS_uniform.trees", printgen=10)
monitors[++mni] = mnModel(filename="output/PS_uniform.log",printgen=10)

# run the analysis
mymcmc = mcmc(mymodel, moves, monitors, nruns=2)
mymcmc.burnin(10000,100)
mymcmc.run(30000)


# summarize output
treetrace = readTreeTrace("output/PS_uniform.trees", treetype="non-clock")
#treetrace.summarize()

map_tree = mapTree(treetrace,"output/PS_uniform_map.tre")


# you may want to quit RevBayes now
q()