################################################################################
#
# RevBayes Example: Bayesian inference of diversification rates under a
#                   episodic birth-death model
#
#
# authors: Sebastian Hoehna
#
################################################################################

#######################
# Reading in the Data #
#######################

### Read in the "observed" tree
T <- readTrees("data/primates_tree.nex")[1]

# Get some useful variables from the data. We need these later on.
taxa <- T.taxa()


# Create some vector for the moves and monitors of this analysis
moves    = VectorMoves()
monitors = VectorMonitors()

NUM_INTERVALS = 100
NUM_BREAKS := NUM_INTERVALS - 1



####################
# Create the rates #
####################

# prior and hyperprior for overall amount of rate variation
speciation_global_scale_hyperprior <- 0.0021
extinction_global_scale_hyperprior <- 0.0021

speciation_global_scale ~ dnHalfCauchy(0,1)
extinction_global_scale ~ dnHalfCauchy(0,1)


# create a random variable at the present time
log_speciation_at_present ~ dnUniform(-10.0,10.0)
log_speciation_at_present.setValue(0.0)
log_extinction_at_present ~ dnUniform(-10.0,10.0)
log_extinction_at_present.setValue(-1.0)

moves.append( mvSlideBactrian(log_speciation_at_present,weight=5))
moves.append( mvSlideBactrian(log_extinction_at_present,weight=5))

for (i in 1:NUM_BREAKS) {
  sigma_speciation[i] ~ dnHalfCauchy(0,1)
  sigma_extinction[i] ~ dnHalfCauchy(0,1)

  # Make sure values initialize to something reasonable
  sigma_speciation[i].setValue(runif(1,0.005,0.1)[1])
  sigma_extinction[i].setValue(runif(1,0.005,0.1)[1])

  # non-centralized parameterization of horseshoe
  delta_log_speciation[i] ~ dnNormal( mean=0, sd=sigma_speciation[i]*speciation_global_scale*speciation_global_scale_hyperprior )
  delta_log_extinction[i] ~ dnNormal( mean=0, sd=sigma_extinction[i]*extinction_global_scale*extinction_global_scale_hyperprior )
}

# Assemble first-order differences and speciation_rate at present into the random field
speciation := fnassembleContinuousMRF(log_speciation_at_present,delta_log_speciation,initialValueIsLogScale=TRUE,order=1)
extinction := fnassembleContinuousMRF(log_extinction_at_present,delta_log_extinction,initialValueIsLogScale=TRUE,order=1)

# Move all field parameters in one go
moves.append( mvEllipticalSliceSamplingSimple(delta_log_speciation,weight=5,tune=FALSE) )
moves.append( mvEllipticalSliceSamplingSimple(delta_log_extinction,weight=5,tune=FALSE) )

# Move all field hyperparameters in one go
moves.append( mvHSRFHyperpriorsGibbs(speciation_global_scale, sigma_speciation , delta_log_speciation , speciation_global_scale_hyperprior, propGlobalOnly=0.75, weight=10) )
moves.append( mvHSRFHyperpriorsGibbs(extinction_global_scale, sigma_extinction , delta_log_extinction , extinction_global_scale_hyperprior, propGlobalOnly=0.75, weight=10) )

# Swap moves to exchange adjacent delta,sigma pairs
moves.append( mvHSRFIntervalSwap(delta_log_speciation ,sigma_speciation ,weight=5) )
moves.append( mvHSRFIntervalSwap(delta_log_extinction ,sigma_extinction ,weight=5) )
interval_times <- abs(T.rootAge() * seq(1, NUM_BREAKS, 1)/NUM_INTERVALS)


### Empirical Taxon Sampling
rho <- 1.0

Galagidae             = clade("Galago_senegalensis",  "Otolemur_crassicaudatus", missing= 17)
Lorisidae             = clade("Perodicticus_potto", "Loris_tardigradus", "Nycticebus_coucang", missing=6)
Cheirogaleoidea       = clade("Cheirogaleus_major", "Microcebus_murinus", missing= 19)
Lemuridae             = clade("Lemur_catta", "Varecia_variegata", missing=17)
Lemuriformes          = clade(Lemuridae, Cheirogaleoidea, missing=29)
Atelidae_Aotidae      = clade("Alouatta_palliata", "Aotus_trivirgatus", missing=30)
NWM                   = clade(Atelidae_Aotidae,"Callicebus_donacophilus", "Saimiri_sciureus", "Cebus_albifrons", missing=93)
Hominoidea            = clade("Pan_paniscus", "Hylobates_lar", missing=19)
Cercopithecoidea      = clade("Colobus_guereza", "Macaca_mulatta", "Chlorocebus_aethiops", missing=60)

missing_species_per_clade = v(Galagidae, Lorisidae, Cheirogaleoidea, Lemuridae, Lemuriformes, Atelidae_Aotidae, NWM, Hominoidea, Cercopithecoidea)

timetree ~ dnEpisodicBirthDeath(rootAge=T.rootAge(), lambdaRates=speciation, lambdaTimes=interval_times, muRates=extinction, muTimes=interval_times, rho=rho, incompleteClades=missing_species_per_clade, condition="survival", taxa=taxa)

### clamp the model with the "observed" tree
timetree.clamp(T)



#############
# The Model #
#############


### workspace model wrapper ###
mymodel = model(rho)

### set up the monitors that will output parameter values to file and screen
monitors.append( mnModel(filename="output_CSrho/primates_EBD.log",printgen=10, separator = TAB) )
monitors.append( mnFile(filename="output_CSrho/primates_EBD_speciation_rates.log",printgen=10, separator = TAB, speciation) )
monitors.append( mnFile(filename="output_CSrho/primates_EBD_speciation_times.log",printgen=10, separator = TAB, interval_times) )
monitors.append( mnFile(filename="output_CSrho/primates_EBD_extinction_rates.log",printgen=10, separator = TAB, extinction) )
monitors.append( mnFile(filename="output_CSrho/primates_EBD_extinction_times.log",printgen=10, separator = TAB, interval_times) )
monitors.append( mnScreen(printgen=1000, extinction_global_scale, speciation_global_scale) )



################
# The Analysis #
################

### workspace mcmc ###
mymcmc = mcmc(mymodel, monitors, moves, nruns=2, combine="mixed")

### run the MCMC ###
mymcmc.run(generations=50000, tuningInterval=200)


## quit ##
q()