# plot traits on tree

setwd("/Users/joelnitta/Dropbox/R/nepenthes/")

library(phytools)
library(ape)
library(gridExtra)
library(grid)
library(ggrepel)
library(ggplot2)
library(plotrix)
library(ips)

source("bin/my_funcs.R")
source("bin/my_themes.R")

####################
### read in data ###
####################

### set input files
trait_file <- "data/BauerData.csv"
moran_data <- "data/WetDrySyndromes.csv"

# STOP: choose and load tree
# tree_select <- select.list(c("ITS", "no_ITS"), title="Select tree")
# 
# if (tree_select == "ITS") {
#   tree_file <- "data/NepITSburninTree.nwk"
# } else if (tree_select == "no_ITS") {
#   tree_file <- "data/BEAST_tree.nwk"
# }
# 
# phy <- read.tree(tree_file)

# New version: read non-ITS tree in nexus file generated by tree annotator, retain 2 signif digits
tree_select <- "no_ITS"
phy <- read.beast("data/Reannotated_Neptree_nonITS", digits=2)

# load traits
traits <- read.csv(trait_file, na.strings = "?")
moran <- read.csv(moran_data, na.strings = "?")

# we have two sources of wax observations. 
# I'm going to keep the wax column from Moran 
# for now because it has fewer NAs
traits$Wax <- NULL

# convert moran StickyFluid to binary
moran$StickyFluid <- as.numeric(moran$StickyFluid)
moran$StickyFluid[which (moran$StickyFluid == 1)] <- 0
moran$StickyFluid[which (moran$StickyFluid == 2)] <- 1

# optional: merge Moran 2013 wet/dry syndrome data with our trait data
rownames(traits) <- traits$Taxon
rownames(moran) <- moran$Taxon
traits <- merge(traits, moran, by = "Taxon")
# merge obliterates rownames, so add again
rownames(traits) <- traits$Taxon

# subsetting
# first make lists of traits
eco_traits <- c("MinAltitude", "MaxAltitude", "AltRange", "Distribution", "MidpointLat", "MidpointLong")
moran_traits <- colnames(moran)
mapping_traits <- c("AltClass", "Diet")
quant_traits <- c("PeristWidth",  "PeristSlope", "LaminaLength", "LowerPitchHeight", "UpperPitchHeight", "LowerPeristRibs", "UpperPeristRibs", "PeristResid")
qual_traits <- c("Dimorphism", "ContractedLowers", "ColorationSimilar", "LowerRedder", "UpperRedder", "ColorDimorphism", "LowerLidContrast", "LowerPeristContrast", "LowerPeristStriped", "UpperLidContrast", "UpperPeristContrast", "UpperPeristStriped", "LaminaIndumentum", "PitcherIndumentum", "Eyespots", "StickyFluid")

# drop ecological traits
traits <- traits[,!(colnames(traits) %in% eco_traits)]
rm(moran)

# choose which coding scheme to use for dimorphism
# single multistate trait "ColorDimorphism": drop other binary traits
traits <- traits[,!(colnames(traits) %in% c("ColorationSimilar", "LowerRedder", "UpperRedder"))]
# OR, keep three binary traits for dimorphism, drop others
# traits <- traits[,!(colnames(traits) == "ColorDimorphism")]

# reset list of qualitative traits to match df (now that we dropped some)
qual_traits <- colnames(traits)[colnames(traits) %in% qual_traits]

###################################
### size-correction and scaling ###
###################################

# to size-correct peristome width (without phylo correction) use residuals
mod1 <- lm(PeristWidth ~ LaminaLength, data = traits)
PeristResid <- residuals(mod1)
PeristResid <- as.data.frame(PeristResid, rownames=names(PeristResid))
traits <- merge(traits, PeristResid, by = "row.names", all=TRUE)
rownames(traits) <- traits$Row.names
traits$Row.names <- NULL
traits$PeristWidth <- NULL
rm(PeristResid)
quant_traits <- quant_traits[quant_traits != "PeristWidth"]

# to size-correct perstiome width (with phylo correction) use phytools
# WORK ON THIS

# scale quantitative traits from 1 to 5, where 1 is min value in that column, and 5 is max
traits[,colnames(traits) %in% quant_traits] <- as.data.frame(apply(as.matrix(traits[,colnames(traits) %in% quant_traits]), MARGIN = 2, rescale, newrange=c(1,5)))

#######################################
### load phylogeny, match to traits ###
#######################################

# set outgroup, root tree
outgroup <- c("Triphyophyllum_peltatum", "Ancistrocladus_abbreviatus")
phy <- root(phy, outgroup)

# ladderize left
phy <- ladderize(phy, right = FALSE)

# trim phylogeny to match trait data
phy <- drop.tip(phy, phy$tip.label[!(phy$tip.label %in% rownames(traits))])

# get traits in same order as tips
traits <- traits[phy$tip.label,]

################
### plotting ###
################

### tree edits:
# abbreviate "Nepenthes" in taxon names to just N.
phy$tip.label <- gsub("Nepenthes_", "N. ", phy$tip.label)

# rescale tree to arbitrary length 100 for easier spacing
phy <- rescaleTree (phy, 100)

# drop posterior probabilities less than 0.5
phy$posterior[phy$posterior < 0.5] <- ""

#### STOP HERE: choose traits to map
map_trait <- select.list(c("Diet", "AltClass"), title="Select taxon trimming scheme")
plot_trait <- select.list(c("Quant", "Qual"), title="Select trait type")
# once traits are chosen, continue

### reproduce ggplot2 color scheme
if (map_trait == "Diet") {
  map.cols <- gg_color_hue(length(levels(traits$Diet))) #make a list of colors, one for each level of mapping trait
  tip.colors <- traits$Diet # copy mapping trait
  levels(tip.colors) <- map.cols # reset levels
  tip.colors <- as.character(tip.colors) # convert to character
} else if (map_trait == "AltClass") {
  map.cols <- gg_color_hue(length(levels(traits$AltClass)))
  tip.colors <- traits$AltClass
  levels(tip.colors) <- map.cols
  tip.colors <- as.character(tip.colors)
}

if (plot_trait == "Quant") {
# print tree with offset taxa labels
pdf(file=make_filename(paste("nepenthes_map_traits", plot_trait, map_trait, tree_select, sep="-"), ".pdf", date=FALSE), width=6.5, height = 9, paper="letter")
spacer <- 3 # need to adjust space between labels if change width of figure
# base plot
plot(phy, 
     no.margin=FALSE, 
     x.lim = 150, # tree is 100 units wide. need to leave this much space for taxa and symbols
     cex=0.4, 
     show.tip.label = TRUE, 
     label.offset=(length(quant_traits)+1)*spacer) # have 7 traits to plot, so taxa lables go in 8th spot
# add colored circles for diet
tiplabels(pch=16, cex=1, col = tip.colors)
# use loop to add circles scaled to traits and trait names
for (i in 1:length(quant_traits)) {
  tiplabels(pch=21,cex=0.25*traits[,which(colnames(traits) == quant_traits[i])], bg = rgb(0,0,0,alpha=0) , col = "black", adj=i*spacer)
  mtext(quant_traits[i], side = 3, at = 100 + i*spacer,srt=60, las=2, cex=0.5, padj=0, adj = 0, line=-1)
}
# add support values at nodes
nodelabels(phy$posterior, adj=c(1.2, -0.3), frame="n", cex=0.5, font=1)
# add legend
legend("topleft", legend = levels(traits[,colnames(traits) == map_trait]), title = map_trait, pch = 16, col = map.cols, cex = 0.8, pt.cex = 1.2, bty = "n")
dev.off()

} else if (plot_trait == "Qual") {

shape <- traits[,which(colnames(traits) %in% qual_traits)]
shape <- apply(shape, 2, function (x) gsub(0, 1, x))
shape <- apply(shape, 2, function (x) gsub(2, 1, x))
shape[is.na(shape)] <- 0
shape <- apply(shape, 2, function (x) as.numeric(x))

fill.col <- traits[,which(colnames(traits) %in% qual_traits)]
fill.col <- apply(fill.col, 2, function(x) gsub(1, "black", x) )
fill.col <- apply(fill.col, 2, function(x) gsub(0, "white", x) )
fill.col <- apply(fill.col, 2, function(x) gsub(2, "grey", x) )

outline.col <- traits[,which(colnames(traits) %in% qual_traits)]
outline.col <- apply(outline.col, 2, function(x) gsub("[0-2]", "darkgrey", x) )

# print tree with offset taxa labels
pdf(file=make_filename(paste("nepenthes_map_traits", plot_trait, map_trait, tree_select, sep="-"), ".pdf", date=FALSE), width=6.5, height = 9, paper="letter")
spacer <- 3 # need to adjust space between labels if change width of figure
# base plot
plot(phy, 
     no.margin=FALSE, 
     x.lim = 150, # tree is 100 units wide. need to leave this much space for taxa and symbols
     cex=0.4, 
     show.tip.label = TRUE, 
     label.offset=(length(qual_traits)+1)*spacer)
# add colored circles for plot trait
tiplabels(pch=16, cex=1, col = tip.colors)
# use loop to add circles scaled to traits and trait names
for (i in 1:length(qual_traits)) {
  tiplabels(pch=22,cex=shape[,i], bg = fill.col[,i], col = outline.col[,i], adj=i*spacer)
  mtext(qual_traits[i], side = 3, at = 100 + i*spacer,srt=60, las=2, cex=0.5, padj=0, adj = 0, line=-1)
}
# add support values at nodes
nodelabels(phy$posterior, adj=c(1.2, -0.3), frame="n", cex=0.5, font=1)
# add legend
legend("topleft", legend = levels(traits[,colnames(traits) == map_trait]), title = map_trait, pch = 16, col = map.cols, cex = 0.8, pt.cex = 1.2, bty = "n")
dev.off()
}