# devtools::install_github("derekmichaelwright/agData")
library(agData)
library(lme4)      # For mixed-effects models
# Prep data
myCaption <- "derekmichaelwright.github.io/dblogr/academic/gxe"

DT::datatable(dd)

# Prep data
xx <- dd %>% select(-VEG, -REP) %>%
  gather(Trait, Value, DTE, DTF, DTS, DTM) %>%
  mutate(Trait = factor(Trait, levels = c("DTE","DTF","DTS","DTM")))
# plot
mp <- ggplot(xx, aes(x = Trait, y = Value, fill = Trait)) +
  geom_boxplot(outlier.shape = NA, alpha = 0.8) +
  facet_wrap(ExptShort ~ ., ncol = 6) +
  scale_fill_manual(values = c("darkgreen","steelblue","purple2","darkorange")) +
  theme_agData(legend.position = "none",
               axis.text.x = element_text(angle = 45, hjust = 1)) +
  labs(title = "Lentil Diversity Panel Phenology Data",
       x = NULL, y = NULL, caption = myCaption)
ggsave("gxe_1_01.png", mp, width = 10, height = 8)

# Prep data
xx <- dd %>% select(-DTE, -DTF, -DTS, -DTM) %>%
  gather(Trait, Value, VEG, REP) %>%
  mutate(Trait = factor(Trait, levels = c("VEG","REP")))
# plot
mp <- ggplot(xx, aes(x = ExptShort, y = Value, fill = Expt)) +
  geom_boxplot(outlier.shape = NA, alpha = 0.8) +
  facet_grid(Trait ~ ., scale = "free_y") +
  scale_fill_manual(values = agData_Colors) +
  theme_agData(legend.position = "none",
               axis.text.x = element_text(angle = 45, hjust = 1)) +
  labs(title = "Lentil Diversity Panel Phenology Data",
       x = NULL, y = NULL, caption = myCaption)
ggsave("gxe_1_02.png", mp, width = 8, height = 6)

# Prep data
xx <- dd %>% gather(Trait, Value, DTE, DTF, DTS, DTM, VEG, REP) %>%
  filter(ExptShort == "Su16") %>%
  mutate(Rep = factor(Rep),
         Trait = factor(Trait, levels = c("DTE","DTF","DTS","DTM","VEG","REP")))
# plot
mp <- ggplot(xx, aes(x = as.factor(Rep), y = Value, fill = Rep)) +
  geom_violin(alpha = 0.8) +
  facet_grid(. ~ Trait) +
  theme_agData(legend.position = "none") +
  labs(title = "Lentil Diversity Panel PHenology Data",
       y = NULL, caption = myCaption)
ggsave("gxe_1_03.png", mp, width = 8, height = 4)

myGxE <- function(xx, myTraits = c("DTE","DTF","DTS","DTM","VEG","REP")) {
  # Initialize an empty list to store variance component data
  var_list <- list()
  blup_list <- list() #store blups values
  # Loop over each trait
  for(i in myTraits) {
    # Define the model formula with ENV as a fixed factor
    formula <- as.formula(paste(i, "~ Expt + (1 | Name) + (1 | Name:Expt) + (1 | Expt:Rep)"))
    # Fit the mixed-effects model
    model <- lmer(formula, data = xx)  # Replace with your actual data
    # Extract variance components for random effects
    var_i <- VarCorr(model) %>% as.data.frame(vc)
    # Calculate total phenotypic variance 
    var_total <- sum(var_i$vcov) 
    # Calculate proportions as percentages
    var_i$Proportion <- (var_i$vcov / var_total) * 100
    #
    var_i <- var_i %>% mutate(Trait = i) %>%
      select(Component=grp, Trait, Proportion, vcov, sdcor)
    # Store results in the list
    var_list[[i]] <- var_i
  }
  # Combine all trait data into one data frame
  var_table <- do.call(rbind, var_list) %>%
    mutate(Trait = factor(Trait, levels = myTraits),
           Component = factor(Component, 
              levels = c("Residual", "Expt:Rep", "Name:Expt", "Name"),
              labels = c("Residual + Env", "Rep:Env", "G x E", "Genetic")))
  # Output
  var_table
}

gg_GxE_bar <- function(xx) {
  myColors <- c("darkred", "purple4", "steelblue", "darkgreen")
  # Create the stacked bar plot
  mp <- ggplot(xx, aes(x = Trait, y = Proportion, fill = Component)) +
    geom_col(position = "stack", color = "black", alpha = 0.7) +
    geom_text(aes(label = round(Proportion)), position = position_stack(vjust = 0.5)) +
    scale_fill_manual(values = myColors) +
    theme_agData_col() +
    labs(title = "Proportion of Phenotypic Variance by Trait",
         y = "Proportion of Total Variance (%)",
         caption = myCaption)
}

gg_GxE_pie <- function(xx, ncol = 3) {
  myColors <- c("darkred", "purple4", "steelblue", "darkgreen")
  # Create the stacked bar plot
  mp <- ggplot(xx, aes(x = "x", y = Proportion, fill = Component)) +
    geom_col(position = "stack", alpha = 0.7) +
    geom_text(aes(label = round(Proportion)), position = position_stack(vjust = 0.5)) +
    coord_polar("y", start = 0) +
    facet_wrap(. ~ Trait, ncol = ncol) +
    scale_fill_manual(values = myColors) +
    scale_x_discrete(limits = c("x_empty", "x")) +
    theme_agData_pie() +
    guides(fill = guide_legend(reverse = T)) +
    labs(title = "Proportion of Phenotypic Variance by Trait",
         y = "Proportion of Total Variance (%)",
         caption = myCaption)
}

# Save variance table 
var_table <- myGxE(dd)
write.csv(var_table, "gxe_All.csv", row.names = F)
# Plot
mp <- gg_GxE_bar(var_table)
ggsave("gxe_All_01.png", mp, width = 6, height = 4)
mp <- gg_GxE_pie(var_table)
ggsave("gxe_All_02.png", mp, width = 6, height = 4)

# Save variance table 
var_table <- myGxE(dd %>% filter(MacroEnv == "Temperate"))
write.csv(var_table, "gxe_Temperate.csv", row.names = F)
# Plot
mp <- gg_GxE_bar(var_table)
ggsave("gxe_Temperate_01.png", mp, width = 6, height = 4)
mp <- gg_GxE_pie(var_table)
ggsave("gxe_Temperate_02.png", mp, width = 6, height = 4)

# Save variance table 
var_table <- myGxE(dd %>% filter(MacroEnv == "Mediterranean"))
write.csv(var_table, "gxe_Mediterranean.csv", row.names = F)
# Plot
mp <- gg_GxE_bar(var_table)
ggsave("gxe_Mediterranean_01.png", mp, width = 6, height = 4)
mp <- gg_GxE_pie(var_table)
ggsave("gxe_Mediterranean_02.png", mp, width = 6, height = 4)

# Save variance table 
var_table <- myGxE(dd %>% filter(MacroEnv == "South Asia"))
write.csv(var_table, "gxe_SouthAsia.csv", row.names = F)
# Plot
mp <- gg_GxE_bar(var_table)
ggsave("gxe_SouthAsia_01.png", mp, width = 6, height = 4)
mp <- gg_GxE_pie(var_table)
ggsave("gxe_SouthAsia_02.png", mp, width = 6, height = 4)