#RAREFACTION FUNCTION
#Reads your data file and plots rarefaction curves for the samples

library(ggplot2)
library(tidyr)
library(ggrepel)

rarefaction <- function(groupnum) {
  
  #Rarefaction function
  rarefact <- function(count_data) {
    Ni <- count_data
    N <- sum(Ni)
    m <- seq(N)
    richness <- sapply(m, function(x) sum(1 - choose(N - Ni, x)/choose(N, x)))
    richness
  }
  
  #Read the data file
  commdata <- read.csv(paste("https://people.ucsc.edu/~mclapham/eart101/papers/group_",groupnum,"_data.csv",sep=""),row.names=1)

  #Perform the rarefaction function on each sample
  rar_res<-apply(commdata, 2, rarefact)
  
  if (class(rar_res) == "matrix") {
    rar_df <- tidyr::gather(as.data.frame(rar_res))
    
    names(rar_df) <- c("sample", "species")
    
    #Generate identification numbers for life and death assemblages
    rar_df$sample <- gsub("\\..*","", rar_df$sample)
    rar_df$specimens <- rep(seq(nrow(as.data.frame(rar_res))), ncol(as.data.frame(rar_res)))
    rar_df$type <- gsub("[0-9]", "", rar_df$sample)
    
    rar_df$label <- ifelse(rar_df$specimens==max(rar_df$specimens), rar_df$sample, NA)
    
    ggplot(rar_df, aes(specimens, species, group=sample)) +
      geom_line(aes(color=type), lwd=1) +
      geom_text_repel(aes(label=label, color=type), show.legend=F, na.rm=T) +
      scale_color_manual(values=c("purple", "orange", "black", "black")) +
      theme_classic()
    
  } else if (class(rar_res) == "list") {
    rar_df <- data.frame("sample" = rep(names(rar_res), sapply(rar_res, length)), species=unlist(rar_res, use.names = F))
    rar_df$sample <- gsub("\\..*","", rar_df$sample)
    rar_df$specimens <- unlist(sapply(rar_res, function(x) seq(length(x))))
    
    rar_df$type <- gsub("[0-9]", "", rar_df$sample)
    
    labels <- data.frame("x" = sapply(rar_res, length),
                         "y" = sapply(rar_res, max),
                         "sample" = names(sapply(rar_res, max)))
    labels$sample <- gsub("\\..*","", labels$sample)
    
    ggplot(rar_df, aes(specimens, species, group=sample)) +
      geom_line(aes(color=type), lwd=1) +
      geom_text_repel(data=labels, aes(x=x, y=y, label=sample), show.legend=F, na.rm=T) +
      scale_color_manual(values=c("purple", "orange", "black", "black")) +
      theme_classic()
  }

  #To save the plot, click the "export" option above the plot window and then select "save as image."
  #Set the file name and use the "directory" button in the save window to choose the folder where it will be saved.  
  
}