logbf1_cutoff = 4
logbf2_cutoff = 4
# Bayes factor cutoffs
# If you want, change these values and re-run the block to see how it affects the results

st_results$target.acc = ifelse(st_results$logbf1 > logbf1_cutoff & st_results$logbf2 > logbf2_cutoff, "Y", "N")
# Label loci with BF1 > 4 and BF2 > 4, this is equivalent to selecting loci
# with "M1" in the best.fit.model column

num_m1 = nrow(subset(st_results, target.acc=="Y"))
# Count the number of M1 loci

bf_spread_st = ggplot(st_results, aes(x=logbf1, y=logbf2, color=target.acc)) +
  geom_point(size=2, alpha=0.25) +
  geom_vline(xintercept=logbf1_cutoff, size=0.75, linetype="dotted", color="#999999") +
  geom_hline(yintercept=logbf2_cutoff, size=0.75, linetype="dotted", color="#999999") +
  ggtitle(paste(num_m1, " loci accelerated in\nmarine mammal lineages (ST)", sep="")) +
  xlab("log BF1") +
  ylab("log BF2") +
  scale_color_manual(values=corecol(pal="wilke", numcol=2)) +
  bartheme() +
  theme(legend.position="none")
bf_spread_st = ggExtra::ggMarginal(bf_spread_st, type="histogram", bins=50, fill=corecol(pal="wilke", numcol=1, offset=4), color="#666666")
# Render the BF figure with ggplot

print(bf_spread_st)

# Display the figure

mm_anc_branches = c("turTru2-orcOrc1", "odoRosDiv1-lepWed1", "triMan1")
# A list of ancestral marine mammal lineages

st_results_m1_m2 = subset(st_results, best.fit.model=="M1" | best.fit.model=="M2")
# Select only loci in which these lineages can be accelerated

st_results_mm_anc = data.frame()
for(i in 1:nrow(st_results_m1_m2)){
  
  cur_model = tolower(st_results_m1_m2[i,]$best.fit.model)


  accel_lineages = st_results_m1_m2[i,][[paste("accel.lineages.", cur_model, sep="")]]
  accel_lineages = strsplit(accel_lineages, split=",")[[1]]
  # Get the accelerated branches from the current locus
  
  accel_anc_lineages = intersect(accel_lineages, mm_anc_branches)
  # Get a list of ancestral marine mammal branches that are accelerated in the current locus
  
  if(setequal(accel_anc_lineages, mm_anc_branches)){
  # Check if the current accelerated ancestral marine mammal branches match the full
  # set of ancestral marine mammal branches
    
    if(all(accel_lineages %in% mm_branches)){
    # Check if all the accelerated lineages in the current locus are marine mammal lineages
      
      st_results_mm_anc = rbind(st_results_mm_anc, st_results_m1_m2[i,])
      # Add this locus to the data frame of ancestrally accelerated marine mammal lineages
    }
  }
}
# Count the number of times all three branches are accelerated and no non-marine mammal 
# lineages are accelerated

st_results_mm_anc %>% kable() %>% kable_styling(bootstrap_options=c("striped", "condended", "responsive"), full_width=F) %>% 
 scroll_box(width = "100%", height = "20vh")

# Display a table with the uniquely accelerated ancestral marine mammal loci

control_branches = c("ailMel1", "loxAfr1", "eleEdw1", "loxAfr3-eleEdw1", "oviAri3", "capHir1", "oviAri3-capHir1", "bosTau8", "bosTau8-oviAri3", "panHod1", "bosTau8-oviAri3", "panHod1-bosTau8")
# A list of control branches sister to marine mammals

tree_info$control = "background"
tree_info[tree_info$label %in% control_branches,]$control = "control"
tree_info[tree_info$label %in% mm_branches,]$control = "marine mammal"
# Label the branches as control, mm, or neither

control_tree = ggtree(tree, size=0.8, ladderize=F, aes(color=tree_info$control)) +
  geom_tiplab(color="#333333", size=4) +
  scale_color_manual(name="Branch types", values=c("background"="#666666", "marine mammal"=corecol(numcol=1), "control"=corecol(numcol=1, offset=5))) +
  geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$control!="background", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  theme(legend.position=c(0.1,0.1))
# Render the tree figure

print(control_tree)

# Display the tree figure

control_anc_branches = c("loxAfr3-eleEdw1", "panHod1-bosTau8", "ailMel1")
# A list of ancestral control branches

st_results_m2 = subset(st_results, best.fit.model=="M2")
# Select only those loci where these branches can be accelerated (M2)

st_results_control_anc = data.frame()
for(i in 1:nrow(st_results_m2)){
  accel_lineages = st_results_m2[i,]$accel.lineages.m2
  accel_lineages = strsplit(accel_lineages, split=",")[[1]]
  # Get the accelerated branches from the current locus
  
  accel_anc_lineages = intersect(accel_lineages, control_anc_branches)
  # Get a list of ancestral control branches that are accelerated in the current locus
  
  if(setequal(accel_anc_lineages, control_anc_branches)){
    # Check if all the accelerated lineages in the current locus are control lineages
    
    st_results_control_anc = rbind(st_results_control_anc, st_results_m2[i,])
    # Add this locus to the data frame of ancestrally accelerated control lineages
  }
}
# Count the number of times all three branches are accelerated and no non-marine mammal 
# lineages are accelerated

st_results_control_anc %>% kable() %>% kable_styling(bootstrap_options=c("striped", "condended", "responsive"), full_width=F) %>% 
 scroll_box(width = "100%", height = "20vh")

# Display a table with the uniquely accelerated ancestral control loci

echo_branches = c("turTru2", "orcOrc1", "turTru2-orcOrc1", "myoLuc2", "myoDav1", "myoDav1-myoLuc2", "eptFus1", "eptFus1-myoDav1")
# A list of echolocating mammals

control_branches = c("capHir1", "oviAri3-capHir1", "bosTau8", "bosTau8-oviAri3", "panHod1", "bosTau8-oviAri3", "panHod1-bosTau8", "pteAle1", "pteVam1", "pteAle1-pteVam1")
# A list of control branches sister to echolocating clades

tree_info$control = "background"
tree_info[tree_info$label %in% control_branches,]$control = "control"
tree_info[tree_info$label %in% echo_branches,]$control = "echolocating"
# Label the branches as control, mm, or neither

control_tree = ggtree(tree, size=0.8, ladderize=F, aes(color=tree_info$control)) +
  geom_tiplab(color="#333333", size=4) +
  scale_color_manual(name="Branch types", values=c("background"="#666666", "echolocating"=corecol(numcol=1), "control"=corecol(numcol=1, offset=5))) +
  geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$control!="background", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  theme(legend.position=c(0.1,0.1))
# Render the tree figure

print(control_tree)

# Display the tree figure

echo_anc_branches = c("turTru2-orcOrc1", "myoDav1-myoLuc2")
# A list of ancestral echolocating lineages

st_results_m1_m2 = subset(st_results, best.fit.model=="M1" | best.fit.model=="M2")
# Select only loci in which these lineages can be accelerated

st_results_echo_anc = data.frame()
for(i in 1:nrow(st_results_m1_m2)){
  cur_model = tolower(st_results_m1_m2[i,]$best.fit.model)

  accel_lineages = st_results_m1_m2[i,][[paste("accel.lineages.", cur_model, sep="")]]
  accel_lineages = strsplit(accel_lineages, split=",")[[1]]
  # Get the accelerated branches from the current locus
  
  accel_anc_lineages = intersect(accel_lineages, echo_anc_branches)
  # Get a list of ancestral echolocating branches that are accelerated in the current locus
  
  if(setequal(accel_anc_lineages, echo_anc_branches)){
  # Check if the current accelerated ancestral echolocating branches match the full
  # set of ancestral mecholocating branches
    
    if(all(accel_lineages %in% echo_branches)){
    # Check if all the accelerated lineages in the current locus are echolocating lineages
      
      st_results_echo_anc = rbind(st_results_echo_anc, st_results_m1_m2[i,])
      # Add this locus to the data frame of ancestrally accelerated echolocating lineages
    }
  }
}
# Count the number of times all three branches are accelerated and no non-echolocating 
# lineages are accelerated

st_results_echo_anc %>% kable() %>% kable_styling(bootstrap_options=c("striped", "condended", "responsive"), full_width=F) %>% 
 scroll_box(width = "100%", height = "20vh")

# Display a table with the uniquely accelerated ancestral marine mammal loci

control_anc_branches = c("panHod1-bosTau8", "pteAle1-pteVam1")
# A list of ancestral control branches

st_results_m2 = subset(st_results, best.fit.model=="M2")
# Select only those loci where these branches can be accelerated (M2)

st_results_control_anc = data.frame()
for(i in 1:nrow(st_results_m2)){
  accel_lineages = st_results_m2[i,]$accel.lineages.m2
  accel_lineages = strsplit(accel_lineages, split=",")[[1]]
  # Get the accelerated branches from the current locus
  
  accel_anc_lineages = intersect(accel_lineages, control_anc_branches)
  # Get a list of ancestral control branches that are accelerated in the current locus
  
  if(setequal(accel_anc_lineages, control_anc_branches)){
    # Check if all the accelerated lineages in the current locus are control lineages
    
    st_results_control_anc = rbind(st_results_control_anc, st_results_m2[i,])
    # Add this locus to the data frame of ancestrally accelerated control lineages
  }
}
# Count the number of times all three branches are accelerated and no non-marine mammal 
# lineages are accelerated

st_results_control_anc %>% kable() %>% kable_styling(bootstrap_options=c("striped", "condended", "responsive"), full_width=F) %>% 
 scroll_box(width = "100%", height = "20vh")

# Display a table with the uniquely accelerated ancestral control loci

h = corecol(numcol=1, pal="wilke", offset=3)
l = corecol(numcol=1, offset=3)
# Colors

scf_tree_fig = ggtree(scf_tree, size=0.8, ladderize=F, aes(color=scf_tree_info$scf)) +
  scale_color_continuous(name='sCF', low=l, high=h, limits=c(0,1)) +
  geom_tiplab(color="#333333", fontface='italic', size=4) +
  theme(legend.position=c(0.1,0.1))
# Render the tree figure

print(scf_tree_fig)

# Display the tree figure

scf_dist = ggplot(data=aln_stats, aes(x=node.scf.avg)) +
  geom_histogram(bins=50, fill=corecol(pal="wilke", numcol=1, offset=1), color="#999999") +
  scale_y_continuous(expand=c(0,0)) +
  xlab("Avg. sCF per locus") +
  ylab("# of loci") +
  bartheme()

print(scf_dist)

scf_data = select(aln_stats, locus, node.scf.avg)
names(scf_data)[1] = "original.id"

all_results = merge(st_results, gt_results, by="original.id", suffixes=c(".st", ".gt"))
all_results = merge(all_results, scf_data, by="original.id")

st_results_high = subset(all_results, node.scf.avg > 0.6)
st_results_high = subset(st_results_high, best.fit.model.st == best.fit.model.gt)
st_results_high = subset(st_results_high, best.fit.model.st != "M0" & best.fit.model.gt != "M0")

high_locus = sample_n(st_results_high, 1)
# This randomly selects a locus with high sCF where the best fit models match between
# gene and species tree methods

high_locus = subset(st_results_high, original.id=="chr5-87")
# Select an M1 locus as an example. Other good ones are: chr5-87

high_st_model = tolower(high_locus$best.fit.model.st)
high_gt_model = tolower(high_locus$best.fit.model.gt)

## Selecting locus with high sCF
##########

st_results_low = subset(all_results, node.scf.avg < 0.45)
st_results_low = subset(st_results_low, best.fit.model.st != best.fit.model.gt)
st_results_low = subset(st_results_low, best.fit.model.st != "M0" & best.fit.model.gt != "M0")

#low_locus = sample_n(st_results_low, 1)
# This randomly selects a locus with low sCF where the best fit models don't match between
# gene and species tree methods

low_locus = subset(st_results_low, original.id=="chr15-72")
# Select an M1 locus as an example. Other good ones are: chr15-72, chr1-11

low_st_model = tolower(low_locus$best.fit.model.st)
low_gt_model = tolower(low_locus$best.fit.model.gt)

## Selecting locus with LOW sCF
##########

con_high_st = strsplit(high_locus[[paste("conserved.lineages.", high_st_model, ".st" , sep="")]], split=",")[[1]]
acc_high_st = strsplit(high_locus[[paste("accel.lineages.", high_st_model, ".st", sep="")]], split=",")[[1]]
# HIGH SCF ST

con_high_gt = strsplit(high_locus[[paste("conserved.lineages.", high_gt_model, ".gt", sep="")]], split=",")[[1]]
acc_high_gt = strsplit(high_locus[[paste("accel.lineages.", high_gt_model, ".gt", sep="")]], split=",")[[1]]
# HIGH SCF GT

con_low_st = strsplit(low_locus[[paste("conserved.lineages.", low_st_model, ".st" , sep="")]], split=",")[[1]]
acc_low_st = strsplit(low_locus[[paste("accel.lineages.", low_st_model, ".st" , sep="")]], split=",")[[1]]
# LOW SCF ST

con_low_gt = strsplit(low_locus[[paste("conserved.lineages.", low_gt_model, ".gt", sep="")]], split=",")[[1]]
acc_low_gt = strsplit(low_locus[[paste("accel.lineages.", low_gt_model, ".gt", sep="")]], split=",")[[1]]
# LOW SCF GT

## Split the lineages in each state for both M1 and M2 into vectors
##########

tree_info$high.state.st = "background"
tree_info[tree_info$label %in% acc_high_st,]$high.state.st = "accelerated"
tree_info[tree_info$label %in% con_high_st,]$high.state.st = "conserved"
# HIGH SCF ST
  
tree_info$high.state.gt = "background"
tree_info[tree_info$label %in% acc_high_gt,]$high.state.gt = "accelerated"
tree_info[tree_info$label %in% con_high_gt,]$high.state.gt = "conserved"
# HIGH SCF GT

tree_info$low.state.st = "background"
tree_info[tree_info$label %in% acc_low_st,]$low.state.st = "accelerated"
tree_info[tree_info$label %in% con_low_st,]$low.state.st = "conserved"
# LOW SCF ST
  
tree_info$low.state.gt = "background"
tree_info[tree_info$label %in% acc_low_gt,]$low.state.gt = "accelerated"
tree_info[tree_info$label %in% con_low_gt,]$low.state.gt = "conserved"
# LOW SCF GT

# Label each branch in the tree as background, accelerated, or conserved for the best fitting model
##########

tree_info$high.bl.st = tree_info$branch.length
tree_info[tree_info$label %in% con_high_st,]$high.bl.st = tree_info[tree_info$label %in% con_high_st,]$high.bl.st * high_locus[[paste("conserved.rate.", low_st_model, ".st" , sep="")]]
tree_info[tree_info$label %in% acc_high_st,]$high.bl.st = tree_info[tree_info$label %in% acc_high_st,]$high.bl.st * high_locus[[paste("accel.rate.", low_st_model, ".st" , sep="")]]
# HIGH SCF ST

tree_info$high.bl.gt = tree_info$branch.length
tree_info[tree_info$label %in% con_high_gt,]$high.bl.gt = tree_info[tree_info$label %in% con_high_gt,]$high.bl.gt * high_locus[[paste("conserved.rate.", low_st_model, ".gt" , sep="")]]
tree_info[tree_info$label %in% acc_high_gt,]$high.bl.gt = tree_info[tree_info$label %in% acc_high_gt,]$high.bl.gt * high_locus[[paste("accel.rate.", low_st_model, ".gt" , sep="")]]
# HIGH SCF GT

tree_info$low.bl.st = tree_info$branch.length
tree_info[tree_info$label %in% con_low_st,]$low.bl.st = tree_info[tree_info$label %in% con_low_st,]$low.bl.st * low_locus[[paste("conserved.rate.", low_st_model, ".st" , sep="")]]
tree_info[tree_info$label %in% acc_low_st,]$low.bl.st = tree_info[tree_info$label %in% acc_low_st,]$low.bl.st * low_locus[[paste("accel.rate.", low_st_model, ".st" , sep="")]]
# LOW SCF ST

tree_info$low.bl.gt = tree_info$branch.length
tree_info[tree_info$label %in% con_low_gt,]$low.bl.gt = tree_info[tree_info$label %in% con_low_gt,]$low.bl.gt * low_locus[[paste("conserved.rate.", low_st_model, ".gt" , sep="")]]
tree_info[tree_info$label %in% acc_low_gt,]$low.bl.gt = tree_info[tree_info$label %in% acc_low_gt,]$low.bl.gt * low_locus[[paste("accel.rate.", low_st_model, ".gt" , sep="")]]
# LOW SCF GT

# M1 LOCUS:  Adjust branch lengths based on the substitution rates in the best fitting model
##########

high_st_tree = tree
high_gt_tree = tree
low_st_tree = tree
low_gt_tree = tree

for(i in 1:length(tree$edge.length)){
  cur_len = tree$edge.length[i]
  cur_len = format(round(cur_len, 8), nsmall=8)
  # Get the original length for this branch and convert it to a rounded string
  
  if(!tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$high.state.st == "background"){
    high_st_tree$edge.length[i] = tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$high.bl.st
  }
  # HIGH SCF ST
  
  if(!tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$high.state.gt == "background"){
    high_gt_tree$edge.length[i] = tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$high.bl.gt
  }
  # HIGH SCF GT

  if(!tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$low.state.st == "background"){
    low_st_tree$edge.length[i] = tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$low.bl.st
  }
  # LOW SCF ST
  
  if(!tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$low.state.gt == "background"){
    low_gt_tree$edge.length[i] = tree_info[format(round(tree_info$branch.length,8), nsmall=8)==cur_len,]$low.bl.gt
  }
  # LOW SCF GT 

}
# Loop over every branch in the tree and adjust the branch length based on the state in the model

####################

cols = c("background"="#666666", "conserved"=corecol(numcol=1, pal="wilke", offset=1), "accelerated"=corecol(numcol=1, pal="wilke"))
# Set colors for conservation states

cur_title = paste("Locus: ", high_locus$original.id, " / Avg. sCF: ", round(high_locus$node.scf.avg,3), "\nMethod: Species tree / Supported model: ", high_locus$best.fit.model.st, sep="")
xmax = max(nodeHeights(high_st_tree)) + sd(nodeHeights(high_st_tree))
high_st_tree_fig = ggtree(high_st_tree, size=0.8, ladderize=F, aes(color=tree_info$high.state.st)) +
  geom_tiplab(color="#333333", size=3.2) +
  xlim(0, xmax) +
  scale_color_manual(name="Rate states", values=cols) +
  #geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$mm=="Y", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  ggtitle(cur_title) +
  theme(legend.position="bottom",
        plot.margin=unit(c(1,1,1,1), "cm"))
# HIGH SCF ST tree

cur_title = paste("Locus: ", high_locus$original.id, " / Avg. sCF: ", round(high_locus$node.scf.avg,3), "\nMethod: Gene tree / Supported model: ", high_locus$best.fit.model.gt, sep="")
xmax = max(nodeHeights(high_gt_tree)) + sd(nodeHeights(high_gt_tree))
high_gt_tree_fig = ggtree(high_gt_tree, size=0.8, ladderize=F, aes(color=tree_info$high.state.gt)) +
  geom_tiplab(color="#333333", size=3.2) +
  xlim(0, xmax) +
  scale_color_manual(name="Rate states", values=cols) +
  #geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$mm=="Y", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  ggtitle(cur_title) +
  theme(legend.position="bottom",
        plot.margin=unit(c(1,1,1,1), "cm"))
# HIGH SCF GT tree

cur_title = paste("Locus: ", low_locus$original.id, " Species tree method / Avg. sCF: ", round(low_locus$node.scf.avg,3), "\nMethod: Species tree / Supported model: ", low_locus$best.fit.model.st, sep="")
xmax = max(nodeHeights(low_st_tree)) + sd(nodeHeights(low_st_tree))
low_st_tree_fig = ggtree(low_st_tree, size=0.8, ladderize=F, aes(color=tree_info$low.state.st)) +
  geom_tiplab(color="#333333", size=3.2) +
  xlim(0, xmax) +
  scale_color_manual(name="Rate states", values=cols) +
  #geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$mm=="Y", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  ggtitle(cur_title) +
  theme(legend.position="bottom",
        plot.margin=unit(c(1,1,1,1), "cm"))
# LOW SCF ST tree

cur_title = paste("Locus: ", low_locus$original.id, " / Avg. sCF: ", round(low_locus$node.scf.avg,3), "\nMethod: Gene tree / Supported model: ", low_locus$best.fit.model.gt, sep="")
xmax = max(nodeHeights(low_gt_tree)) + sd(nodeHeights(low_gt_tree))
low_gt_tree_fig = ggtree(low_gt_tree, size=0.8, ladderize=F, aes(color=tree_info$low.state.gt)) +
  geom_tiplab(color="#333333", size=3.2) +
  xlim(0, xmax) +
  scale_color_manual(name="Rate states", values=cols) +
  #geom_text(aes(x=branch, label=ifelse(tree_info$node.type=="internal" & tree_info$mm=="Y", tree_info$label, '')), size=3, vjust=-0.3, show.legend=F) +
  ggtitle(cur_title) +
  theme(legend.position="bottom",
        plot.margin=unit(c(1,1,1,1), "cm"))
# LOW SCF GT tree

tree_legend = get_legend(high_st_tree_fig)
# Get the legend from one fig

tree_fig_top = plot_grid(high_st_tree_fig + theme(legend.position="none"),
                     high_gt_tree_fig + theme(legend.position="none"),
                     low_st_tree_fig + theme(legend.position="none"),
                     low_gt_tree_fig + theme(legend.position="none"),
                     ncol=2)
# Combine the M1 locus trees into a single row, removing individual legends

tree_fig = plot_grid(tree_fig_top, tree_legend, nrow=2, rel_heights=c(1,0.1))
# Combine the M1 locus row, M2 locus row, and a single legend

print(tree_fig)

# Display the figure