---
title: "Figure2"
author: "Meredith Hanlon"
date: "2023-04-11"
output: html_document
---

```{r global_options, echo = FALSE}
knitr::opts_chunk$set(warning=FALSE, message=FALSE, echo=FALSE, error = TRUE)
```

```{r}
library(broom)
library(tidyverse)
library(rstatix)
library(ggpubr)
library(readxl)
library(stats)
library(viridis)
library(patchwork)
library(rmarkdown)
library(ggrepel)

stderror <- function(x) sd(x, na.rm = TRUE) / sqrt(length(na.omit(x)))

theme_mh <- function(){ 
    font <- "sans"   #assign font family up front
    theme(
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      panel.background = element_blank(),
      panel.border = element_rect(color = "black", fill = NA),
      axis.ticks = element_line(color = "black"),   
      axis.line = element_line(color = "black"),
      strip.background = element_blank(),
      legend.key = element_rect(fill = NA),
      text = element_text(family = font, size = 10, color = "black"),
      axis.text = element_text(family = font, size = 8, color = "black")
    )
}
```

# Figure 2A

```{r}
lxf <- read_excel("Hanlon_2023_data.xlsx", sheet = "Figure_2A_1")
icf <- read_excel("Hanlon_2023_data.xlsx", sheet = "Figure_2A_2")

pca_fit <- lxf %>% 
  dplyr::select(ends_with("ear")) %>% 
  select(Al_ear:Sn_ear) %>% 
  prcomp(center = TRUE, scale = TRUE)

# get variations
summary(pca_fit)
# PC1 = 33.6, PC2 = 14.4

rots_pre <- pca_fit %>%
  tidy(matrix = "rotation") %>% 
  filter(PC <= 2) 
rots_top <- rots_pre %>% 
  separate(column, into = c("Element", "ear_id"), sep = "_") %>% 
  ungroup() %>% 
  group_by(Element) %>% 
  slice_max(abs(value), n = 1) %>% 
  unite(column, Element, ear_id, sep = "_") %>% 
  dplyr::select(column, value) 
rots <- rots_pre %>% 
  filter(column %in% rots_top$column) %>% 
  separate(column, into = c("Element", "ear_id"), sep = "_") 

pcd <- pca_fit %>% augment(lxf)
rd <- rots %>% 
    # need to scale rots to match
  mutate(value = value * 15) %>% 
  pivot_wider(names_from = "PC", names_prefix = "PC", values_from = "value") 

arrow_style <- arrow(angle = 20, ends = "first", type = "closed", length = grid::unit(8, "pt"))  

bips1 <- ggplot() +
  geom_segment(data = rd, aes(PC1, PC2), xend = 0, yend = 0, arrow = arrow_style, color = "grey60") +
  geom_point(data = pcd, aes(.fittedPC1, .fittedPC2, color = Genotype, shape = Week), size = 3) +
  scale_color_manual(values = c("#A8641B", "#F59D40", "#0A7EA8", "#39C4F6")) +
  geom_text_repel(data = rd, aes(PC1, PC2, label = Element), color = "black", box.padding = 0.1) +
  xlab("PC 1 (33.6%)") +
  ylab("PC 2 (14.4%)") +
  theme_mh() +
  theme(legend.position = "top",
        legend.box = "vertical",
        legend.margin = margin(t = -5, b = -5))

pca_fit <- icf %>% 
  dplyr::select(ends_with("ear")) %>% 
  select(Al_ear:Cr_ear) %>% 
  prcomp(center = TRUE, scale = TRUE)

# get variations
summary(pca_fit)
# PC1 = 30.8, PC2 = 14.3

arrow_style <- arrow(angle = 20, ends = "first", type = "closed", length = grid::unit(8, "pt"))

rots_pre <- pca_fit %>%
  tidy(matrix = "rotation") %>% 
  filter(PC <= 2) 

rots_top <- rots_pre %>% 
  separate(column, into = c("Element", "ear_id"), sep = "_") %>% 
  ungroup() %>% 
  group_by(Element) %>% 
  slice_max(abs(value), n = 1) %>% 
  unite(column, Element, ear_id, sep = "_") %>% 
  dplyr::select(column, value) 

rots <- rots_pre %>% 
  filter(column %in% rots_top$column) %>% 
  separate(column, into = c("Element", "ear_id"), sep = "_") 

pcd <- pca_fit %>% augment(icf)
rd <- rots %>% 
    # need to scale rots to match biplot
  mutate(value = value * 15) %>% 
  pivot_wider(names_from = "PC", names_prefix = "PC", values_from = "value") 

bips3 <- ggplot() +
  geom_segment(data = rd, aes(PC1, PC2), xend = 0, yend = 0, arrow = arrow_style, color = "grey60") +
  geom_point(data = pcd, aes(.fittedPC1, .fittedPC2, color = Genotype, shape = Week), size = 3) +
  scale_color_manual(values = c("#A8641B", "#F59D40", "#0A7EA8", "#39C4F6")) +
  geom_text_repel(data = rd, aes(PC1, PC2, label = Element), color = "black", box.padding = 0.1) +
  xlab("PC 1 (30.8%)") +
  ylab("PC 2 (14.3%)") +
  theme_mh() +
  theme(legend.position = "top",
        legend.box = "vertical",
        legend.margin = margin(t = -5, b = -5))

f2av2 <-   bips1 + ggtitle("XRF") + 
  theme(plot.title = element_text(family = "sans", size = 12, hjust = 0.5)) +
  bips3 + ggtitle("ICP") +
  theme(plot.title = element_text(family = "sans", size = 12, hjust = 0.5)) +
  plot_layout(width = c(4, 4), nrow = 1, ncol = 2, guides = "collect") & theme(legend.position = 'top')
f2av2

ggsave("Figure2A.png", f2av2, width = 8, height = 4, units = "in")
```

# Figure 2B

Summary of PCAs - all PCAs were run and organized using this code

runs were changed to include either all leaves, ear leaves, or three (ear/young/old leaves)

This was done across all plants inidividually or as plot averages

Data were either centered, scaled, both, or neither

```{r}
pca_fit <- lxf %>% 
  select(Al_old:Sn_ear8) %>% # changed as appropriate
  prcomp(center = FALSE, scale = FALSE) #changed as appropriate

pmpc <- pca_fit %>% 
  augment(lxf) %>% 
  dplyr::select(Week, Genotype, pm_pc1 = .fittedPC2, total_root_length:d95) 

models <- pmpc %>% 
  dplyr::select(Week, Genotype, pm_pc1, total_root_length:d95) %>% # changed to incldue PC1, PC2, or PC1 + PC2
  pivot_longer(total_root_length:d95, names_to = "trait", values_to = "trait_value") %>% 
  dplyr::select(-Week, -Genotype) %>% 
  ungroup() %>% 
  nest(data = -trait) %>% 
  mutate(fit = purrr::map(data, ~lm(trait_value ~ ., data = .x)),
         tidied = purrr::map(fit, tidy),
    glanced = purrr::map(fit, glance),
    augmented = purrr::map(fit, augment)) 

models %>% 
  unnest(glanced) %>% 
  paged_table()

m1 <- models %>% unnest(glanced) %>% 
  dplyr::select(trait, r.squared:nobs)

# data were output into excel, organized and used for plotting below
```


```{r}
pcrs <- read_xlsx("Hanlon_2023_data.xlsx", sheet = "Figure_2B")

pcrs3 <- pcrs %>% 
  dplyr::select(trait, adjr2 = adj.r.squared, p = p.value, predic, method, leaves, center, scale, grouping) %>% 
  filter(center == "T" & scale == "T") %>% 
  filter(trait %in% c("total_root_length", "shallow_20", "deep_30", "deep_40", "d50", "d70", "d95"))

pcrs3$trait <- factor(pcrs3$trait, 
         levels = c("total_root_length", "shallow_20", "deep_30", "deep_40", "d50", "d70", "d95"),
         labels = c("Total", "<20cm",">30cm", ">40cm", "d50", "d70", "d95"))

pcrs3$grouping <- factor(pcrs3$grouping, levels = c("ind", "plot"), labels = c("Individual plants", "Plot average"))

pcrs3 %>% 
  unite(meth2, leaves, method, sep = ", ") %>% 
  mutate(adjr2 = case_when(adjr2 >= 0 ~ adjr2, TRUE ~ 0)) %>% 
  mutate(sig = case_when(p < 0.05 ~ "sig1",
                         p >= 0.05 & p < 0.1 ~ "sig2")) %>% 
  mutate(r3 = case_when(sig %in% c("sig1", "sig2") ~ adjr2)) %>% 
  ggplot(aes(x = predic, y = trait, fill = adjr2)) +
  geom_tile(color = "white", linewidth = 0.1) +
  scale_fill_viridis(limits = c(0,1), expand = c(0,0), name = expression(R^2), na.value = "white") +
  geom_text(aes(label = round(r3, 2), color = sig),  size = 3) +
  scale_color_manual(values = c("black", "white", NA), guide = "none") +
  theme_mh() +
  scale_y_discrete(limits = rev) +
  xlab(NULL) +
  facet_grid(grouping ~ meth2) +
  theme(legend.position = "top",
        panel.spacing.x = unit(0, "lines"),
        legend.margin = margin(0, 0, 0, 0),
        legend.box.margin = margin(b = -15),
        axis.ticks.x = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank(),
        legend.key.height = unit(2, "mm"),
        legend.title = element_text(margin = margin(b = 10, t = -2))) 

ggsave("Figure2B.png", width = 8, height = 4, units = "in")
```

