####################################
#Functions to support stomata analyses#
####################################
#Process images for scoring
stomataImProcesseR <- function(folderID){
  #Infer path based on folder ID
  path <- paste0(getwd(), "/", folderID)
  
  #List all images in path
  imID <- list.files(path, pattern=".jpg")
  
  for(i in 1:length(imID)){
    #Load image
    imraw <- magick::image_read(paste0(path, "/", imID[i]))
    
    #Crop image (a lot of rotating and choping on the right end of image)
    imprep <- magick::image_rotate(magick::image_chop(magick::image_rotate(imraw, 90), "70"),-90)
    imprep <- magick::image_rotate(magick::image_chop(magick::image_rotate(imprep, 180), "198"), -180)
 
    #Add cell IDs
    #As
    imprep <- magick::image_annotate(imprep, "A1", size = 50, color = "red", location = "+10+10")
    imprep <- magick::image_annotate(imprep, "A2", size = 50, color = "red", location = "+10+230")
    imprep <- magick::image_annotate(imprep, "A3", size = 50, color = "red", location = "+10+450")
    
    #Bs
    imprep <- magick::image_annotate(imprep, "B1", size = 50, color = "red", location = "+229+10")
    imprep <- magick::image_annotate(imprep, "B2", size = 50, color = "red", location = "+229+230")
    imprep <- magick::image_annotate(imprep, "B3", size = 50, color = "red", location = "+229+450")
    
    #Cs
    imprep <- magick::image_annotate(imprep, "C1", size = 50, color = "red", location = "+439+10")
    imprep <- magick::image_annotate(imprep, "C2", size = 50, color = "red", location = "+439+230")
    imprep <- magick::image_annotate(imprep, "C3", size = 50, color = "red", location = "+439+450")
    
    #Ds
    imprep <- magick::image_annotate(imprep, "D1", size = 50, color = "red", location = "+659+10")
    imprep <- magick::image_annotate(imprep, "D2", size = 50, color = "red", location = "+659+230")
    imprep <- magick::image_annotate(imprep, "D3", size = 50, color = "red", location = "+659+450")
    
    #Es
    imprep <- magick::image_annotate(imprep, "E1", size = 50, color = "red", location = "+879+10")
    imprep <- magick::image_annotate(imprep, "E2", size = 50, color = "red", location = "+879+230")
    imprep <- magick::image_annotate(imprep, "E3", size = 50, color = "red", location = "+879+450")
    
    #Save images in Images_to_score
    #Check if folder exists
    # if not then creates it
    output_dir <- file.path(paste0(path, "/Images_to_score"))
    
    if(dir.exists(output_dir)){
      #print(paste0("Dir", output_dir, " already exists!"))
    }else{
      print(paste0("Created ", output_dir))
      dir.create(output_dir)  
    }
    
    magick::image_write(imprep, path = paste0(output_dir, "/To_score_", imID[i]))
  }
  
  #Move all raw images in right Raw_Images folder
  #Check if folder exists
  # if not then creates it
  output_dir <- file.path(paste0(path, "/Raw_Images/"))
  
  if(dir.exists(output_dir)){
    print(paste0("Dir", output_dir, " already exists!"))
  }else{
    print(paste0("Created ", output_dir))
    dir.create(output_dir)  
  }
  
  for(i in 1:length(imID)){
    file.rename(from = paste0(path, "/", imID[i]), to= paste0(path, "/Raw_Images/", imID[i]))
  }
  
  print("All images to score are in Images_to_score/ and \n raw images are in Raw_Images/")
}

#Arguments
#folderID <- SHOULD BE Images_to_score/
#Outputs:
# DB1 and DB2

stomataDBMakeR <- function(folderID){
  #Infer path based on folder ID
  path <- paste0(getwd(), "/", folderID)
  
  #List all images in path
  imID <- list.files(path, pattern='.jpg')
  
  ###
  #Create & populate DB1
  ###
  folder <- gsub("/Images_to_score", "", folderID)
  DB1 <- matrix(ncol=8, nrow=length(imID))
  colnames(DB1) <- c("Plant_ID", "Surface", "Position", "Image_ID","N_stomata", "N_trichomes", "Avg_stomata_length", "Avg_stomata_width")
  DB1 <- as.data.frame(DB1)
  
  #Populate DB1
  DB1$Plant_ID <- rep(folder, nrow(DB1))
  DB1$Image_ID <- imID
  DB1$Surface <- c(rep("T", 3), rep("B", 3))
  DB1$Position <- rep(seq(from=1, to=3), 2)
  
  #Write DB1
  path2 <- gsub("/Images_to_score", "", path)
  print(paste("Create DB1: ", paste0("DB1_", folder, ".csv")))
  write.csv(DB1, file = paste(path2, paste0("/DB1_", folder, ".csv"), sep=''), row.names = F, quote = F)
  
  ###
  #Create & populate DB2 for scoring
  ###
  
  #Define cells for scoring stomata and trichomes
  cells <- paste0(sort(rep(LETTERS[seq(from = 1, to = 5)], 3)), rep(seq(1,3), 3))
  
  DB2 <- NULL
  for(i in 1:nrow(DB1)){
    tmp <- cbind(as.vector(rep(DB1$Image_ID[i], length(cells))), cells)
    DB2 <- rbind(DB2, tmp)
  }
  
  #Convert to DF
  DB2 <- as.data.frame(DB2)
  colnames(DB2) <- c("Image_ID", "Cell")
  
  #Add 2 cols
  DB2$N_stomata <- NA
  DB2$N_trichomes <- NA
  
  
  #Write DB2
  print(paste("Create DB2: ", paste0("DB2_", folder, ".csv")))
  write.csv(DB2, file = paste(path2, paste0("/DB2_", folder, ".csv"), sep=''), row.names = F, quote = F)
  
  #Tell user to
  print(paste("Enter # stomata and trichomes per cell in", paste0("DB2_", folder, ".csv")))
} 

#######
#UPDATE DB1 with # stomata and trichomes from DB2
#######

#Argument
#folderID <- "G1_b3_1_6"
#Output
# Updates DB1 with data from DB2
stomataDB1UpdateR1 <- function(folderID){
  #Infer path based on folder ID
  path <- paste0(getwd(), "/", folderID)
  
  #Load DB1 & DB2
  DB1 <- read.csv(paste(path, paste0("/DB1_", folderID, ".csv"), sep=''))
  DB2 <- read.csv(paste(path, paste0("/DB2_", folderID, ".csv"), sep=''))
  
  for(i in 1:nrow(DB1)){
    #Subset DB2 based on DB1
    tmp <- subset(DB2, DB2$Image_ID == DB1$Image_ID[i])
    #Count N stomata and trichomes and update DB1
    DB1$N_stomata[i] <- colSums(tmp[3])
    DB1$N_trichomes[i] <- colSums(tmp[4])
  }
  
  #Write updated DB1
  print(paste("Update DB1: ", paste0("DB1_", folderID, ".csv")))
  write.csv(DB1, file = paste(path, paste0("/DB1_", folderID, ".csv"), sep=''), row.names = F, quote = F)
}

#Input: folderID: Path to DB2
#       nstomata: N stomata measured in pict (1 per cell)
stomataMeasureR <- function(folderID, nstomata){
  #Infer path based on folder ID
  path <- paste0(getwd(), "/", folderID)
  
  #Open DB2
  DB2 <- read.csv(paste(path, paste0("/DB2_", folderID, ".csv"), sep=''))
  
  #List of images
  ImID <- as.vector(unique(DB2$Image_ID))
  
  #Create DB3
  DB3 <- NULL
  for(i in 1:length(ImID)){
    #Subset by ImageID and cells with >= 1 stomata
    tmp <- subset(DB2, DB2$Image_ID == ImID[i] & DB2$N_stomata > 0)
    #Randomly sample N cells equal to nstomata
    tmp <- subset(tmp, tmp$Cell %in% as.vector(sample(tmp$Cell, nstomata)))
    DB3 <- rbind(DB3, tmp[,1:2])
  }
  DB3$Stomate_Length <- NA
  DB3$Stomate_Width <- NA
  DB3$GuardCell_Length <- NA
  DB3$GuardCell_Width <- NA
  DB3$Aperture_Length <- NA
  DB3$Aperture_Width <- NA
  
  #Write out DB3
  print(paste("Create DB3: ", paste0("DB3_", folderID, ".csv")))
  write.csv(DB3, file = paste(path, paste0("/DB3_", folderID, ".csv"), sep=''), row.names = F, quote = F)
  
  #Tell user to
  print(paste("Measure stomata in", paste0("DB3_", folderID, ".csv")))
  
}

#folderID <- "G1_b3_1_6"
#Output
# Updates DB1 with data from DB3
stomataDB1UpdateR2 <- function(folderID){
  #Infer path based on folder ID
  path <- paste0(getwd(), "/", folderID)
  
  #Open DB1/DB3
  DB1 <- read.csv(paste(path, paste0("/DB1_", folderID, ".csv"), sep=''))
  
  DB3 <- read.csv(paste(path, paste0("/DB3_", folderID, ".csv"), sep=''))
  
  for(i in 1:nrow(DB1)){
    tmp <- subset(DB3, DB3$Image_ID == DB1$Image_ID[i])
    #Update DB1
    DB1$Avg_stomata_length[i] <- mean(tmp$Stomate_Length)
    DB1$Avg_stomata_width[i] <- mean(tmp$Stomate_Width)
  }
  #Write updated DB1
  print(paste("Update DB1: ", paste0("DB1_", folderID, ".csv")))
  write.csv(DB1, file = paste(path, paste0("/DB1_", folderID, ".csv"), sep=''), row.names = F, quote = F)
}