#Function to prune the tree, predict from the tree, and get plots

treeFun <- function(data, toPlot = T, title = ""){
  
  #Create Complex Tree - or full tree
  myTree <- tree(Y ~ ., data = data, model = T)
  
  #Perform CV on tree object
  cvTree <- cv.tree(myTree)
  optTree <- which.min(cvTree$dev)
  bestTree <- cvTree$size[optTree]
  
  #Prune tree based on CV results
  pruneTree <- prune.tree(myTree, best = bestTree)
  
  #Use tree to predict original data
  treePred <- predict(pruneTree)
  treeResid <- resid(pruneTree)
  myRange <- range(treePred, data$Y)
  
  #If plotting is selected
  if(toPlot){
    
    #Plot Unpruned Tree
    plot(myTree)
    text(myTree, cex = .75)
    title(main = paste("Unpruned Tree for", title))
    
    #Plot CV
    plot(cvTree$size, cvTree$dev, type = "b", 
         main = paste("Cross Validation for", title))
    
    #Plot Prunned Tree
    plot(pruneTree)
    text(pruneTree, cex = .75)
    title(main = paste("Pruned Tree for", title))
    
    #Plot Observed vs Predicted
    plot(data$Y, treePred, xlim = myRange, ylim = myRange, xlab="Observed TOC Concentration (mg/L) ", ylab="Modeled TOC Concentration (mg/L)",
         main = paste("Pred vs True for", title))
    lines(data$Y, data$Y, col = "black")
    
  } 
  return(pruneTree)
}