library(scales)
library(ggplot2)
library(gridExtra)
library(tidyr)

options(stringsAsFactors = F)

source_dist <- 0.5 # fraction of protonated source
uptake_rate <- 1000 # uptake of A
Vol = 1e7
time_steps <- 2000

# step 1, A -> B
Km_1 <- 1e-2
Kcat_1 <- 10000

# step 2, B -> C
Km_2 <- 1e-2
Kcat_2 <- 10000

# step 3, C is removed and accumulate overtime
C_removal_rate <- c(0.01)

parameter_sets <- data.frame(name = c("No KIE", "Weak KIE (R1:1.5, R2:2)", "Strong KIE (R1:3, R2:5)"),
                             KIE_1 = c(1, 1.5, 3),
                             KIE_2 = c(1, 2, 5))

molecules_list <- list()
reactions_list <- list()

for(a_param_set in 1:nrow(parameter_sets)){
  
  KIE_1 <- parameter_sets$KIE_1[a_param_set]
  KIE_2 <- parameter_sets$KIE_2[a_param_set]
  
  Apool <- c(1000, 0)
  Bpool <- c(1000, 0)
  Cpool <- c(1000, 0)
  Csink <- c(0, 0)
  
  Atrack <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps))
  Btrack <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps))
  Ctrack <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps))
  
  Csink_track <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps))
  
  Rxn_1_track <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps), T = rep(NA, time_steps))
  Rxn_2_track <- data.frame(U = rep(NA, time_steps), L = rep(NA, time_steps), T = rep(NA, time_steps))
  
  
  Atrack[1,] <- Apool
  Btrack[1,] <- Bpool
  Ctrack[1,] <- Cpool
  Csink_track[1,] <- Csink
  
  for(t in 2:time_steps){
    
    # Uptake
    new_A_labelled <- rbinom(1, uptake_rate, source_dist)
    Apool[1] <- Apool[1] + new_A_labelled
    Apool[2] <- Apool[2] + (uptake_rate - new_A_labelled)
    Atrack[t,] <- Apool
    
    # Step 1
    new_B <- c(rpois(1, Kcat_1 * (Apool[1]/Vol) / ((Apool[1]/Vol) + Km_1)), 
               rpois(1, Kcat_1 / KIE_1 * (Apool[2]/Vol) / ((Apool[2]/Vol) + Km_1)))
    Apool <- Apool - new_B # A removed
    Bpool <- Bpool + new_B # B added
    Btrack[t,] <- Bpool
    Rxn_1_track[t,] <- c(new_B, sum(new_B))
    
    # Step 2
    new_C <- c(rpois(1, Kcat_2 * (Bpool[1]/Vol) / ((Bpool[1]/Vol) + Km_2)), 
               rpois(1, Kcat_2 / KIE_2 * (Bpool[2]/Vol) / ((Bpool[2]/Vol) + Km_2)))
    Bpool <- Bpool - new_C # A removed
    Cpool <- Cpool + new_C
    Ctrack[t,] <- Cpool
    Rxn_2_track[t,] <- c(new_C, sum(new_C))
    
    # Step 3
    molecules_removed <- rpois(1, C_removal_rate * sum(Cpool))
    if(molecules_removed != 0){
      labelled_molecules <- rbinom(1, molecules_removed, Cpool[2]/sum(Cpool))
      
      Csink[1] <- Csink[1] + (molecules_removed - labelled_molecules)
      Csink[2] <- Csink[2] + labelled_molecules
      
      Cpool <- Cpool - c((molecules_removed - labelled_molecules), labelled_molecules)
      Csink_track[t,] <- Csink
      
    }else{
      Csink_track[t,] <- Csink
    }
  } 
  
  Afrac <- Atrack[,2] / rowSums(Atrack)
  Bfrac <- Btrack[,2] / rowSums(Btrack)
  Cfrac <- Ctrack[,2] / rowSums(Ctrack)
  Csink_frac <- Csink_track[,2] / rowSums(Csink_track)
  
  parameter_sets$name[a_param_set]
  
  molecules_list[[a_param_set]] <- data.frame(simulation = parameter_sets$name[a_param_set], 
                                              rbind(data.frame(specie = "A", time = 1:time_steps, frac = Afrac),
                                                    data.frame(specie = "B", time = 1:time_steps, frac = Bfrac),
                                                    data.frame(specie = "C", time = 1:time_steps, frac = Cfrac)))
  
  reactions_list[[a_param_set]] <- data.frame(simulation = parameter_sets$name[a_param_set],
                                              rbind(data.frame(reaction = "R1", time = 1:time_steps, Rxn_1_track),
                                                    data.frame(reaction = "R2", time = 1:time_steps, Rxn_2_track)))
  
}

molecules_list <- do.call("rbind", molecules_list)
reactions_list <- do.call("rbind", reactions_list)
reactions_list <- gather(reactions_list, "Labeling form", "Molecules", -(simulation:time))

theme_empty <- theme_bw() +
  theme(line = element_blank(),
        rect = element_blank(),
        strip.text = element_blank(),
        axis.text = element_blank(),
        plot.title = element_blank(),
        axis.title = element_blank(),
        plot.margin = structure(c(0, 0, -1, -1), unit = "lines", valid.unit = 3L, class = "unit"))

species_layout_coord <- data.frame(label = c("A", "B", "C"), x = 0, y = c(0, -2, -4))
reaction_arrows <- data.frame(x = 0, xend = 0, y = c(1.5, -0.5, -2.5, -4.5), yend = c(0.5, -1.5, -3.5, -5.5))
reaction_names <- data.frame(label = c("Input", "R1", "R2", "Removal"), x = 1.2, y = c(1, -1, -3, -5))

schema_plot <- ggplot() +
  geom_point(data = species_layout_coord, aes(x = x, y = y), size = 20, color = "cadetblue2") +
  geom_text(data = species_layout_coord, aes(x = x, y = y, label = label), size = 10) +
  geom_segment(data = reaction_arrows, aes(x = x, xend = xend, y = y, yend = yend), arrow = arrow(), color = "coral1", size = 2) +
  geom_text(data = reaction_names, aes(x = x, y = y, label = label), size = 10) +
  theme_empty + coord_equal() + expand_limits(x = c(-1, 2.5))


scatter_facet_theme <- theme(text = element_text(size = 23),
                             panel.background = element_rect(fill = "gray93"),
                             legend.position = "top",
                             panel.grid.minor = element_blank(),
                             panel.grid.major = element_line(size = 0.5, color = "gray50"),
                             axis.text.x = element_text(angle = 90),
                             axis.text = element_text(color = "black"),
                             axis.ticks = element_line(size = 1),
                             panel.margin = unit(1, "lines"),
                             axis.title.y=element_text(vjust=0.2))

species_plot <- ggplot(molecules_list, aes(x = time, y = frac, color = simulation, group = simulation)) +
  facet_grid(specie ~ .) + 
  geom_hline(yintercept = 0.5, size = 1) +
  geom_hline(yintercept = 0, size = 2) + geom_vline(xintercept = 0, size = 2) +
  geom_line(size = 1) +
  scale_color_brewer(palette = "Set1") + 
  scale_x_log10("Time steps elapsed", breaks = c(10, 100, 1000), expand = c(0,0)) +
  scale_y_continuous("Labeling percentage", labels = percent_format(), expand = c(0,0)) + expand_limits(y = c(0,1)) +
  scatter_facet_theme + theme(strip.background = element_rect(fill = "cadetblue2"))

legend = gtable::gtable_filter(ggplot_gtable(ggplot_build(species_plot)), "guide-box") 
species_plot <- species_plot + theme(legend.position="none")

flux_plot <- ggplot(reactions_list[reactions_list$`Labeling form` != "T",], aes(x = time, y = Molecules, color = simulation, group = simulation)) +
  facet_grid(reaction ~ `Labeling form`) + 
  geom_hline(yintercept = 0.5, size = 1) +
  geom_hline(yintercept = 0, size = 2) + geom_vline(xintercept = 0, size = 2) +
  geom_line(size = 0.6, alpha = 0.6) +
  scale_color_brewer(palette = "Set1", guide = "none") + 
  scale_x_log10("Time steps elapsed", breaks = c(10, 100, 1000), expand = c(0,0)) +
  scale_y_continuous("Molecules converted / time step", expand = c(0,0)) +
  scatter_facet_theme + theme(strip.background = element_rect(fill = "coral1"))

plot_list <- list(schema_plot, species_plot, flux_plot)

core_plots <- marrangeGrob(plot_list, nrow=1, ncol=3, widths = c(1, 2, 2), top=NULL)

jpeg(file = "KIE_example.jpeg", height = 8, width = 14, units = "in", res = 1024)
do.call(grid.arrange, c(list(legend), core_plots, list(heights = c(1,9))))
dev.off()