# =============================================================================
# 4. Analyze metagenomics data from biobakery output using R (WITH REPLICATES)
#    Project: Soil Metagenomics 2026 — Loc1 vs Loc4 comparison (Pseudo-replicates)
# =============================================================================

# -- Prerequisites -----------------------------------------------------------
# install.packages(c("phyloseq", "ggplot2", "vegan", "dplyr",
#                     "tidyr", "pheatmap", "patchwork", "openxlsx", "tibble"))
# devtools::install_github("biobakery/Maaslin2")

# -- Load packages -----------------------------------------------------------
library(phyloseq)
library(ggplot2)
library(vegan)
library(dplyr)
library(tidyr)
library(pheatmap)
library(Maaslin2)
library(tibble)
library(openxlsx)

# -- Working directory -------------------------------------------------------
setwd("~/DATA/Data_Tam_Metagenomics_2026_Soil/reports_rep/")
dir.create("figures", showWarnings = FALSE)
dir.create("tables",  showWarnings = FALSE)


# =================================================================================
# PART 1: DESCRIPTIVE COMPARISONS & DATA EXPORTS
# =================================================================================

# =============================================================================
# STEP 1 — Load MetaPhlAn species-level profiles
# =============================================================================
mpa_file <- "../biobakery_input_rep/results/metaphlan/merged/metaphlan_taxonomic_profiles.tsv"
mpa_raw  <- read.table(mpa_file, header = TRUE, sep = "\t",
                       comment.char = "", check.names = FALSE)
colnames(mpa_raw)[1] <- "clade_name"
rownames(mpa_raw)    <- mpa_raw$clade_name
mpa_otu              <- mpa_raw[, -1, drop = FALSE]

# Keep species-level rows only (s__) and tidy names
species_idx   <- grep("s__", rownames(mpa_otu))
species_otu   <- mpa_otu[species_idx, , drop = FALSE]
species_names <- gsub(".*s__", "", rownames(species_otu))
species_names <- gsub("\\|.*", "", species_names)
species_names <- gsub("_", " ", species_names)
rownames(species_otu) <- species_names

# =============================================================================
# STEP 2 — Build metadata (Updated for Pseudo-Replicates)
# =============================================================================
sample_names <- colnames(species_otu)
sample_names_clean <- gsub("_taxonomic_profile$", "", sample_names)
colnames(species_otu) <- sample_names_clean

# Extract Location and Replicate from names like "Soil_Loc1_rep1"
locations <- case_when(
  grepl("Loc1", sample_names_clean, ignore.case = TRUE) ~ "Loc1",
  grepl("Loc4", sample_names_clean, ignore.case = TRUE) ~ "Loc4",
  TRUE ~ "Unknown"
)

replicates <- case_when(
  grepl("rep1", sample_names_clean, ignore.case = TRUE) ~ "rep1",
  grepl("rep2", sample_names_clean, ignore.case = TRUE) ~ "rep2",
  TRUE ~ "rep_unknown"
)

metadata <- data.frame(
  row.names = sample_names_clean,
  Location  = factor(locations, levels = c("Loc1", "Loc4")),
  Replicate = replicates,
  stringsAsFactors = FALSE
)

cat("Metadata:\n")
print(metadata)

# =============================================================================
# STEP 3 — Create phyloseq object (species)
# =============================================================================
tax_mat <- matrix(rownames(species_otu), ncol = 1,
                  dimnames = list(rownames(species_otu), "Species"))

species_ps <- phyloseq(
  otu_table(species_otu, taxa_are_rows = TRUE),
  sample_data(metadata),
  tax_table(tax_mat)
)

# =============================================================================
# STEP 4 — Load HUMAnN pathway abundances
# =============================================================================
path_file <- "../biobakery_input_rep/results/humann/merged/pathabundance_relab.tsv"
path_raw <- read.table(path_file, header = TRUE, sep = "\t",
                       comment.char = "", check.names = FALSE,
                       row.names = 1, quote = "")

colnames(path_raw) <- gsub("_Abundance$", "", colnames(path_raw))

pathway_names <- rownames(path_raw)
pathway_names <- gsub("UniRef90_", "", pathway_names)
pathway_names <- gsub("_", " ", pathway_names)
rownames(path_raw) <- pathway_names

# Subset & reorder columns to match metadata
path_otu <- path_raw[, sample_names_clean, drop = FALSE]

path_tax_mat <- matrix(rownames(path_otu), ncol = 1,
                       dimnames = list(rownames(path_otu), "Pathway"))

pathway_ps <- phyloseq(
  otu_table(path_otu, taxa_are_rows = TRUE),
  sample_data(metadata),
  tax_table(path_tax_mat)
)

cat("\nSample name verification:\n")
cat("  Match:", all(sample_names(species_ps) == sample_names(pathway_ps)), "\n")


# =============================================================================
# STEP 5 — Visualisations (PNG) + Complete Excel exports
# =============================================================================

# Helper: safe log2 fold change
safe_log2fc <- function(x, y, pseudo = 1e-6) {
  log2((x + pseudo) / (y + pseudo))
}

# --- 5a. Top-N species stacked bar plot ---
top_n_sp <- 20
top_species <- names(sort(rowMeans(otu_table(species_ps)), decreasing = TRUE))[1:top_n_sp]
ps_top <- prune_taxa(top_species, species_ps)

df_species <- psmelt(ps_top) %>%
  mutate(Species = factor(Species, levels = rev(top_species)))

p_species <- ggplot(df_species, aes(x = Sample, y = Abundance, fill = Species)) +
  geom_bar(stat = "identity", position = "fill", width = 0.8) +
  facet_grid(~Location, scales = "free_x", space = "free_x") +
  coord_flip() +
  scale_fill_viridis_d(option = "D") +
  labs(title = paste("Top", top_n_sp, "Species by Sample (Faceted by Location)"),
       x = "Sample", y = "Relative Abundance", fill = "Species") +
  theme_minimal(base_size = 11) +
  theme(legend.position = "bottom", axis.text.x = element_text(angle = 45, hjust = 1))

ggsave("figures/species_top20_barplot.png", p_species, width = 12, height = 8, dpi = 300, bg = "white")

# --- 5b. Species heatmap ---
otu_mat <- as.matrix(otu_table(species_ps))
keep <- apply(otu_mat, 1, max) > 0.0001
otu_filt <- otu_mat[keep, , drop = FALSE]

ann_col <- data.frame(Location = metadata[colnames(otu_filt), "Location"],
                      row.names = colnames(otu_filt))

png("figures/species_heatmap.png", width = 8, height = max(8, 0.25 * nrow(otu_filt) + 2), units = "in", res = 300)
pheatmap(otu_filt, scale = "row", annotation_col = ann_col,
         main = "Species Abundance Heatmap (row-scaled)",
         fontsize_row = 6, fontsize_col = 10, show_rownames = nrow(otu_filt) <= 80)
dev.off()

# --- 5c & 5d. Pathway plots ---
top_n_pw <- 20
top_pw_names <- names(sort(rowMeans(otu_table(pathway_ps)), decreasing = TRUE))[1:top_n_pw]
ps_pw_top <- prune_taxa(top_pw_names, pathway_ps)

df_pw <- psmelt(ps_pw_top) %>%
  mutate(Pathway = factor(Pathway, levels = rev(top_pw_names)))

p_pw <- ggplot(df_pw, aes(x = Sample, y = Abundance, fill = Pathway)) +
  geom_bar(stat = "identity", position = "fill", width = 0.8) +
  facet_grid(~Location, scales = "free_x", space = "free_x") +
  coord_flip() +
  scale_fill_viridis_d(option = "C") +
  labs(title = paste("Top", top_n_pw, "HUMAnN Pathways by Sample")) +
  theme_minimal(base_size = 11) +
  theme(legend.position = "bottom", axis.text.x = element_text(angle = 45, hjust = 1))

ggsave("figures/pathways_top20_barplot.png", p_pw, width = 12, height = 8, dpi = 300, bg = "white")


# --- 5e & 5f. Export complete lists to Excel (Dynamic for Replicates) ---
loc1_cols <- grep("Loc1", colnames(otu_table(species_ps)), value = TRUE)
loc4_cols <- grep("Loc4", colnames(otu_table(species_ps)), value = TRUE)

# SPECIES EXPORT
sp <- as.data.frame(otu_table(species_ps))
sp$Species <- rownames(sp)
sp$Mean_Loc1 <- rowMeans(sp[, loc1_cols, drop = FALSE])
sp$Mean_Loc4 <- rowMeans(sp[, loc4_cols, drop = FALSE])
sp$Diff      <- sp$Mean_Loc4 - sp$Mean_Loc1
sp$Log2FC    <- safe_log2fc(sp$Mean_Loc4, sp$Mean_Loc1)
sp$In_Loc1   <- sp$Mean_Loc1 > 0
sp$In_Loc4   <- sp$Mean_Loc4 > 0
sp$Total_Mean <- sp$Mean_Loc1 + sp$Mean_Loc4

cols_to_keep <- c("Species", loc1_cols, loc4_cols, "Mean_Loc1", "Mean_Loc4", "Diff", "Log2FC", "In_Loc1", "In_Loc4", "Total_Mean")
sp <- sp[, cols_to_keep]
sp <- sp[order(-abs(sp$Diff)), ]

wb1 <- createWorkbook()
addWorksheet(wb1, "All"); writeData(wb1, "All", sp)
addWorksheet(wb1, "Top50_Diff"); writeData(wb1, "Top50_Diff", head(sp, 50))
addWorksheet(wb1, "Top50_Abund"); writeData(wb1, "Top50_Abund", sp[order(-sp$Total_Mean), ][1:50, ])
addWorksheet(wb1, "Loc1_only"); writeData(wb1, "Loc1_only", sp[sp$In_Loc1 & !sp$In_Loc4, ])
addWorksheet(wb1, "Loc4_only"); writeData(wb1, "Loc4_only", sp[!sp$In_Loc1 & sp$In_Loc4, ])
addWorksheet(wb1, "Shared"); writeData(wb1, "Shared", sp[sp$In_Loc1 & sp$In_Loc4, ])

hs <- createStyle(textDecoration = "bold", bgFill = "#D3D3D3")
for (sh in sheets(wb1)) { addStyle(wb1, sh, style = hs, rows = 1, cols = 1:ncol(sp), gridExpand = TRUE); freezePane(wb1, sh, firstRow = TRUE) }
saveWorkbook(wb1, "tables/species_Loc1_vs_Loc4.xlsx", overwrite = TRUE)

# PATHWAYS EXPORT (Same logic)
pw <- as.data.frame(otu_table(pathway_ps))
pw$Pathway <- rownames(pw)
pw$Mean_Loc1 <- rowMeans(pw[, loc1_cols, drop = FALSE])
pw$Mean_Loc4 <- rowMeans(pw[, loc4_cols, drop = FALSE])
pw$Diff      <- pw$Mean_Loc4 - pw$Mean_Loc1
pw$Log2FC    <- safe_log2fc(pw$Mean_Loc4, pw$Mean_Loc1)
pw$In_Loc1   <- pw$Mean_Loc1 > 0
pw$In_Loc4   <- pw$Mean_Loc4 > 0
pw$Total_Mean <- pw$Mean_Loc1 + pw$Mean_Loc4

pw_cols <- c("Pathway", loc1_cols, loc4_cols, "Mean_Loc1", "Mean_Loc4", "Diff", "Log2FC", "In_Loc1", "In_Loc4", "Total_Mean")
pw <- pw[, pw_cols]
pw <- pw[order(-abs(pw$Diff)), ]

wb2 <- createWorkbook()
addWorksheet(wb2, "All"); writeData(wb2, "All", pw)
addWorksheet(wb2, "Top50_Diff"); writeData(wb2, "Top50_Diff", head(pw, 50))
for (sh in sheets(wb2)) { addStyle(wb2, sh, style = hs, rows = 1, cols = 1:ncol(pw), gridExpand = TRUE); freezePane(wb2, sh, firstRow = TRUE) }
saveWorkbook(wb2, "tables/pathways_Loc1_vs_Loc4.xlsx", overwrite = TRUE)


# =================================================================================
# PART 2: DIFFERENTIAL ABUNDANCE ANALYSIS (MaAsLin2)
# =================================================================================

# Prepare data for MaAsLin2
# MaAsLin2 expects: input_data (rows=features, cols=samples) and
# input_metadata (rows=samples, cols=variables) with matching row/col names.

species_df <- as.data.frame(otu_table(species_ps))

# FIX: Create metadata_df and explicitly set row names to match sample IDs
metadata_df <- data.frame(sample_id = rownames(metadata), metadata)
rownames(metadata_df) <- rownames(metadata)

# 2.1 MaAsLin2 for Species
cat("\n--- Running MaAsLin2 for Species ---\n")
fit_species <- Maaslin2(
  input_data = species_df,
  input_metadata = metadata_df,
  output = "Maaslin2_species_output",
  fixed_effects = c("Location"),
  reference = c("Location,Loc1"), # Loc1 is reference, testing Loc4 vs Loc1
  random_effects = c(),
  correction = "BH",
  standardize = FALSE,
  min_abundance = 0.0001,
  min_prevalence = 0.1,
  transform = "NONE"
)

sig_species <- subset(fit_species$results, qval < 0.05 & metadata == "Location" & value == "Loc4")
if(nrow(sig_species) > 0) {
  cat("=== Significantly Different Species (FDR < 0.05) ===\n")
  print(sig_species[, c("feature", "value", "coef", "pval", "qval")] %>% dplyr::arrange(qval))
} else {
  cat("No significantly different species found at FDR < 0.05.\n")
}
write.csv(fit_species$results, "Maaslin2_species_full_results.csv", row.names=FALSE)

# 2.2 MaAsLin2 for Pathways
cat("\n--- Running MaAsLin2 for Pathways ---\n")
pathway_df <- as.data.frame(otu_table(pathway_ps))

fit_pathways <- Maaslin2(
  input_data = pathway_df,
  input_metadata = metadata_df,
  output = "Maaslin2_pathway_output",
  fixed_effects = c("Location"),
  reference = c("Location,Loc1"),
  random_effects = c(),
  correction = "BH",
  standardize = FALSE,
  min_abundance = 0.0001,
  min_prevalence = 0.1,
  transform = "NONE"
)

sig_pathways <- subset(fit_pathways$results, qval < 0.05 & metadata == "Location" & value == "Loc4")
if(nrow(sig_pathways) > 0) {
  cat("=== Significantly Different Pathways (FDR < 0.05) ===\n")
  print(sig_pathways[, c("feature", "value", "coef", "pval", "qval")] %>% dplyr::arrange(qval) %>% head(20))
} else {
  cat("No significantly different pathways found at FDR < 0.05.\n")
}
write.csv(fit_pathways$results, "Maaslin2_pathway_full_results.csv", row.names=FALSE)


# =================================================================================
# PART 3: COMMUNITY STRUCTURE (PERMANOVA)
# =================================================================================

dist_bc_species <- vegdist(t(otu_table(species_ps)), method="bray")

# Test main effect of Location
adonis_species <- adonis2(dist_bc_species ~ Location, data=metadata, permutations=999)

cat("\n=== PERMANOVA Results for Species (Location Effect) ===\n")
print(adonis_species)

write.csv(adonis_species, "permanova_species_results.csv")


# =================================================================================
# PART 4: DIVERSITY VISUALIZATIONS
# =================================================================================

# 4.1 Alpha Diversity (Shannon)
shannon_div <- estimate_richness(species_ps, measures="Shannon")
plot_alpha <- data.frame(
  Sample = rownames(metadata),
  Location = metadata$Location,
  Replicate = metadata$Replicate,
  Shannon = shannon_div$Shannon
)

p_alpha <- ggplot(plot_alpha, aes(x=Location, y=Shannon, fill=Location)) +
  geom_boxplot(alpha=0.5, width=0.5, outlier.shape = NA) +
  geom_point(size=3, position=position_jitter(width=0.1)) +
  theme_minimal() +
  labs(title="Alpha Diversity (Shannon Index)", subtitle="Loc1 vs Loc4 (Pseudo-replicates shown as points)") +
  theme(legend.position = "none")

ggsave("figures/alpha_diversity_shannon.png", p_alpha, width=8, height=6, dpi=300, bg="white")

# 4.2 Beta Diversity (PCoA Plot)
pcoa <- cmdscale(dist_bc_species, k=2, eig=TRUE)
pcoa_df <- data.frame(
  Sample = rownames(metadata),
  Location = metadata$Location,
  Axis1 = pcoa$points[,1],
  Axis2 = pcoa$points[,2]
)

p_pcoa <- ggplot(pcoa_df, aes(x=Axis1, y=Axis2, color=Location, shape=Location)) +
  geom_point(size=4) +
  theme_minimal() +
  labs(title="PCoA of Species Composition (Bray-Curtis)",
       subtitle="Pseudo-replicates should cluster tightly by Location") +
  theme(legend.position = "right")

ggsave("figures/pcoa_species.png", p_pcoa, width=8, height=6, dpi=300, bg="white")

# 4.3 Beta Diversity Boxplot (Within vs Between)
dist_mat <- as.matrix(dist_bc_species)
dist_df <- expand.grid(Sample1 = rownames(dist_mat), Sample2 = colnames(dist_mat))
dist_df$Distance <- as.vector(dist_mat)
dist_df <- dist_df[dist_df$Sample1 != dist_df$Sample2, ]

dist_df$Loc1 <- metadata[dist_df$Sample1, "Location"]
dist_df$Loc2 <- metadata[dist_df$Sample2, "Location"]
dist_df$Comparison <- ifelse(dist_df$Loc1 == dist_df$Loc2, "Within-Location", "Between-Location")

p_beta <- ggplot(dist_df, aes(x=Comparison, y=Distance, fill=Comparison)) +
  geom_boxplot(alpha=0.7) +
  geom_jitter(width=0.1, alpha=0.5) +
  theme_minimal() +
  labs(title="Beta Diversity (Bray-Curtis Distance)",
       subtitle="Within-Location vs Between-Location Distances") +
  theme(legend.position = "none")

ggsave("figures/beta_diversity_boxplot.png", p_beta, width=8, height=6, dpi=300, bg="white")

cat("\n=== Analysis Complete ===\n")
cat("All figures saved to 'figures/' and tables to 'tables/'.\n")