############################################################
# Mitochondrial Morphology Analysis
#
# DESCRIPTION:
#   This script performs morphological analysis and subtype
#   classification of mitochondria from 3D fluorescence images
#   segmented in ZEISS arivis Pro. It includes:
#     - Data import and metadata annotation
#     - Per-sample aggregation and group comparison
#     - Morphometric feature derivation
#     - Unsupervised clustering (k-means)
#     - UMAP dimensionality reduction
#     - Random Forest subtype classification
#     - Statistical testing and visualization
#
# INPUT:
#   - One .xlsx file per biological sample, exported from arivis Pro
#   - All files must be located in the same folder
#
# OUTPUT:
#   - Box plots of volume, surface area, and sphericity per group
#   - UMAP plots colored by subtype and experimental group
#   - Feature importance plot from Random Forest
#   - Subtype proportion plots per group
#   - CSV files with full results and statistics
#   - PNG/PDF figures
#
# AUTHORS: Wayne Mitchell, Cecília G. de Magalhães, Praju Vikas Anekal, Vadim N. Gladyshev
# DATE:    9th April 2026
# VERSION: 2
############################################################

# ── Dependencies ──────────────────────────────────────────────────────────────
required_packages <- c(
  "readxl", "dplyr", "tidyr", "purrr", "ggplot2", "ggpubr", "patchwork",
  "scales", "randomForest", "caret", "umap", "factoextra",
  "RColorBrewer", "mclust", "outliers", "tidyverse"
)
installed <- rownames(installed.packages())
for (pkg in required_packages) {
  if (!pkg %in% installed) install.packages(pkg)
}
invisible(lapply(required_packages, library, character.only = TRUE))


################################################################################
## USER CONFIGURATION                                                         
##                                                                            
## Edit this section to adapt the script to your dataset. All paths,         
## sample names, group assignments, and column names of your dataframe.   


# ── File paths ────────────────────────────────────────────────────────────────
input_path  <- "/add_your_path_here/"    # folder with .xlsx files
output_path <- "/add_your_path_here/" # folder for all outputs

# ── Sample-to-group metadata ──────────────────────────────────────────────────
# List every sample ID (after simplification) and its experimental group.
# Add or remove rows to match your experiment — any number of samples and
# any number of groups are supported.
metadata <- data.frame(
  source_file = c(
    "Sample-01", "Sample-02", "Sample-03", "Sample-04", "Sample-05",
    "Sample-06", "Sample-07", "Sample-08", "Sample-09", "Sample-10",
    "Sample-11", "Sample-12", "Sample-13", "Sample-14", "Sample-25",
    "Sample-16", "Sample-17", "Sample-18", "Sample-19", "Sample-20"
  ),
  Group = c(
    rep("Control", 10),
    rep("Treated", 10)
  )
)

# ── Column name mapping ───────────────────────────────────────────────────────
# Map the exact column headers in your .xlsx files to the internal names
# used throughout this script. Change the LEFT side only based on the software column names output.
# Here we used our arivis pipeline output.
column_map <- c(
  Volume         = "Volume, Volume (µm³)",
  VoxelCount     = "VoxelCount, Volume",
  RelativeVolume = "RelativeVolume, Volume",
  SA_Voxel       = "Surface Area (Voxel), Surface Area (µm²)",
  SA_Mesh        = "Surface Area (Mesh), Surface Area (µm²)",
  Sph_Voxel      = "Sphericity (Voxel), Sphericity",
  Sph_Mesh       = "Sphericity (Mesh), Sphericity",
  Vol_Mesh       = "Volume (Mesh), Sphericity (µm³)"
)


################################################################################
## 1. IMPORTING AND ORGANIZING DATA                                                       


# Read all .xlsx files. 
# From the ZEISS arivis software pipeline you would have one excel file per sample.
xlsx_files <- list.files(
  path = input_path, pattern = "\\.xlsx$",
  full.names = TRUE, recursive = FALSE
)
stopifnot("No .xlsx files found in input_path" = length(xlsx_files) > 0)

data_list        <- lapply(xlsx_files, read_excel)
names(data_list) <- tools::file_path_sans_ext(basename(xlsx_files))
combined_df      <- bind_rows(data_list, .id = "source_file") #here we combine all samples dataframes

# Attach group metadata
combined_df <- left_join(combined_df, metadata, by = "source_file")

# Rename columns and compute derived morphometric features
morpho_df <- combined_df %>%
  rename(any_of(setNames(column_map, names(column_map)))) %>%
  mutate(
    SA_to_Vol       = SA_Mesh / Volume,
    Equiv_diameter  = 2 * ((3 * Volume) / (4 * pi))^(1/3),
    Compactness     = (SA_Mesh^3) / (36 * pi * Volume^2),
    Vol_discrepancy = Vol_Mesh  / Volume,
    SA_discrepancy  = SA_Mesh   / SA_Voxel,
    Sph_discrepancy = Sph_Mesh  / Sph_Voxel
  ) %>%
  filter(if_all(where(is.numeric), is.finite))

message("Total particles loaded: ", nrow(morpho_df))


################################################################################
## 2. FILTERING                                                               
##                                                                            
## Rationale: Segmentation pipelines routinely produce sub-resolution noise  
## objects that are not biologically meaningful. These appear as a pronounced 
## left-tail spike in the log-volume distribution, well below the size of    
## the smallest observed mitochondrial fragments. At the upper end, objects  
## > 10 µm³ likely represent over segmented mitochondrial networks           
## or imaging artefacts.                                                      
##                                                                            
## Absolute cutoffs are preferred over percentile-based trimming because     
## they apply a consistent biological criterion across all samples,           
## irrespective of sample-level differences in particle count or size.        
##                                                                            
## Threshold selection:                                                       
##  1. Visual inspection of volume histograms to identify noise               
##  2. Check particles retained after application of candidate thresholds;   
##     check for the interval with the steepest count drop.                   
##  3. Confirm the threshold visually in the volume distribution plot.       
##                                        
##  Here, our final thresholds: lower = 0.003 µm³, upper = 10 µm³.             


# Diagnostic: volume distribution for representative samples
# Edit samples_to_check to inspect different samples before committing to a threshold
samples_to_check <- c("Sample-01", "Sample-02", "Sample-03", "Sample-04", "Sample-17", "Sample-18", "Sample-19", "Sample-20")

df_check <- morpho_df %>%
  filter(source_file %in% samples_to_check) %>%
  mutate(source_file = factor(source_file, levels = samples_to_check))

quantile_lines <- df_check %>%
  group_by(source_file) %>%
  summarise(
    q05 = quantile(Volume, 0.05), q10 = quantile(Volume, 0.10),
    q20 = quantile(Volume, 0.20), q80 = quantile(Volume, 0.80),
    q90 = quantile(Volume, 0.90), q95 = quantile(Volume, 0.95),
    .groups = "drop"
  ) %>%
  tidyr::pivot_longer(-source_file, names_to = "quantile", values_to = "Volume")

p_quant <- ggplot(df_check, aes(x = log(Volume))) +
  geom_histogram(aes(fill = Group), bins = 60, alpha = 0.7,
                 color = "white", linewidth = 0.2) +
  geom_vline(data = quantile_lines,
             aes(xintercept = log(Volume), color = quantile, linetype = quantile),
             linewidth = 0.7) +
  geom_vline(xintercept = log(vol_min), color = "darkgreen",
             linewidth = 1.2, linetype = "solid") +
  geom_vline(xintercept = log(vol_max), color = "black",
             linewidth = 1.2, linetype = "solid") +
  annotate("text", x = log(vol_min) + 0.15, y = Inf,
           label = paste0(vol_min, " µm³"), vjust = 2, hjust = 0,
           size = 3, color = "darkgreen") +
  annotate("text", x = log(vol_max) - 0.15, y = Inf,
           label = paste0(vol_max, " µm³"), vjust = 2, hjust = 1,
           size = 3, color = "black") +
  scale_color_manual(values = c(q05 = "#e07b39", q10 = "#c0392b", q20 = "#8e44ad",
                                q80 = "#8e44ad", q90 = "#c0392b", q95 = "#e07b39")) +
  scale_linetype_manual(values = c(q05 = "dotted", q10 = "dashed", q20 = "solid",
                                   q80 = "solid",  q90 = "dashed", q95 = "dotted")) +
  scale_x_continuous(breaks = seq(-10, 4, by = 1),
                     minor_breaks = seq(-10, 4, by = 0.5)) +
  facet_wrap(~ source_file, ncol = 2, scales = "free_y") +
  scale_fill_manual(values = group_colors) +
  theme_classic() +
  theme(strip.text      = element_text(face = "bold", size = 12),
        axis.text.x     = element_text(size = 9, angle = 45, hjust = 1),
        legend.position = "top",
        panel.grid.minor.x = element_line(color = "grey90", linewidth = 0.3)) +
  labs(title = "Volume distribution with filter thresholds",
       x = "log(Volume) (µm³)", y = "Number of particles",
       fill = "Group", color = "Quantile", linetype = "Quantile")

print(p_quant)
ggsave(file.path(output_path, "filter_threshold_diagnostic.png"),
       p_quant, width = 10, height = 7, dpi = 300)

# Retention table across candidate thresholds
cat("\n── Particle retention across candidate lower bounds ──\n")
purrr::map_dfr(c(0.0009, 0.003, 0.005, 0.007, 0.01), function(thr) {
  morpho_df %>%
    filter(Volume >= thr, Volume <= vol_max) %>%
    group_by(source_file, Group) %>%
    summarise(n = n(), .groups = "drop") %>%
    mutate(threshold = thr)
}) %>%
  group_by(threshold, Group) %>%
  summarise(mean_particles = round(mean(n)), .groups = "drop") %>%
  tidyr::pivot_wider(names_from = Group, values_from = mean_particles) %>%
  print()

# Particles outside [vol_min, vol_max] are excluded as segmentation artifacts. 
# CHANGE THE VALUES BASED ON YOUR RESULTS
vol_min <- 0.003   # lower bound: removes sub-resolution noise
vol_max <- 10      # upper bound: removes large artifactual object
# Apply final filter
morpho_df_filtered <- morpho_df %>%
  filter(Volume >= vol_min, Volume <= vol_max)

# Per-sample retention summary
retention_summary <- morpho_df %>%
  group_by(source_file, Group) %>%
  summarise(
    total    = n(),
    retained = sum(Volume >= vol_min & Volume <= vol_max),
    pct_kept = round(100 * retained / total, 1),
    .groups  = "drop"
  ) %>%
  arrange(Group, source_file)

cat("\n── Per-sample retention after filtering ──\n")
print(retention_summary)
message("Particles after filtering: ", nrow(morpho_df_filtered),
        " (", round(100 * nrow(morpho_df_filtered) / nrow(morpho_df), 1), "% retained)")


################################################################################
## 3. PLOTTING VOLUME DISTRIBUTION AND MORPHOLOGY BOXPLOTS                   


# Shared theme, choose your colors
#define your gorup colors for the next plots
group_colors <- c("Control" = "#4E9BB9", "Treated" = "firebrick")
dot_colors   <- c("Control" = "#2a7a9b", "Treated" = "firebrick4")

my_theme <- theme_classic() +
  theme(
    axis.title       = element_text(size = 18, color = "grey20"),
    axis.text        = element_text(size = 18, color = "grey20"),
    axis.line        = element_line(color = "grey25"),
    axis.ticks       = element_line(color = "grey25"),
    legend.position  = "none",
    plot.title       = element_text(size = 20, face = "bold", hjust = 0.5, color = "grey15"),
    panel.background = element_rect(fill = "white"),
    plot.background  = element_rect(fill = "white", color = NA)
  )

comparisons <- combn(unique(metadata$Group), 2, simplify = FALSE)

# Per-sample median summary
# Individual mitochondria within a sample are not independent observations;
# between-group comparisons are performed at the sample level using
# per-sample medians (n = number of biological replicates).
df_summary <- morpho_df_filtered %>%
  group_by(source_file, Group) %>%
  summarise(
    Volume      = median(Volume,    na.rm = TRUE),
    SurfaceArea = median(SA_Voxel,  na.rm = TRUE),
    Sphericity  = median(Sph_Voxel, na.rm = TRUE),
    .groups = "drop"
  )

# Helper to avoid repeating boxplot code
make_boxplot <- function(df, y_var, y_label, title) {
  ggplot(df, aes(x = Group, y = .data[[y_var]], fill = Group)) +
    geom_boxplot(width = 0.45, outlier.shape = NA, alpha = 1,
                 color = "grey20", linewidth = 1) +
    geom_jitter(aes(color = Group), width = 0.1, size = 3, alpha = 0.8, shape = 16) +
    stat_compare_means(comparisons = comparisons, method = "wilcox.test",
                       label = "p.signif", size = 5, tip.length = 0.01) +
    scale_fill_manual(values  = group_colors) +
    scale_color_manual(values = dot_colors) +
    labs(title = title, y = y_label, x = "") +
    my_theme
}

p_vol <- make_boxplot(df_summary, "Volume",      "Median Volume (µm³)",       "Volume")
p_sa  <- make_boxplot(df_summary, "SurfaceArea", "Median Surface Area (µm²)", "Surface Area")
p_sph <- make_boxplot(df_summary, "Sphericity",  "Median Sphericity",         "Sphericity")

combined_plot <- p_vol | p_sa | p_sph
print(combined_plot)
ggsave(file.path(output_path, "morphology_boxplots.pdf"),
       combined_plot, width = 12, height = 5)
ggsave(file.path(output_path, "morphology_boxplots.png"),
       combined_plot, width = 12, height = 5, dpi = 300)

# Volume distribution per sample (boxplot)
sample_order <- morpho_df_filtered %>%
  select(source_file, Group) %>%
  distinct() %>%
  arrange(Group, source_file) %>%
  pull(source_file)

p_vol_per_sample <- ggplot(
  morpho_df_filtered,
  aes(x = factor(source_file, levels = sample_order), y = log(Volume), fill = Group)
) +
  geom_boxplot(outlier.shape = NA, alpha = 0.7, color = "grey30", linewidth = 0.6) +
  stat_summary(fun = mean, geom = "point", shape = 23, size = 2.5,
               fill = "white", color = "grey20") +
  scale_fill_manual(values = group_colors) +
  scale_x_discrete(guide = guide_axis(angle = 45)) +
  theme_classic() +
  theme(axis.title = element_text(size = 13), axis.text = element_text(size = 11),
        plot.title = element_text(size = 14, face = "bold", hjust = 0.5),
        legend.position = "top") +
  labs(title = "Mitochondrial volume distribution per sample",
       x = "Sample", y = "log(Volume) (µm³)", fill = "Group")

print(p_vol_per_sample)
ggsave(file.path(output_path, "volume_distribution_per_sample.png"),
       p_vol_per_sample, width = 12, height = 5, dpi = 300)

# Pooled volume distribution by group (relative frequency)
p_vol_distribution <- ggplot(
  morpho_df_filtered,
  aes(x = Volume, fill = Group, color = Group)
) +
  geom_histogram(
    aes(y = after_stat(count / tapply(count, group, sum)[group])),
    bins = 80, alpha = 0.4, position = "identity", linewidth = 0.2
  ) +
  scale_fill_manual(values  = group_colors) +
  scale_color_manual(values = dot_colors) +
  scale_x_continuous(breaks = seq(-10, 4, by = 1)) +
  scale_y_continuous(labels = percent_format(accuracy = 0.1)) +
  theme_classic() +
  theme(axis.title = element_text(size = 13), axis.text = element_text(size = 11),
        axis.text.x = element_text(angle = 45, hjust = 1),
        plot.title  = element_text(size = 14, face = "bold", hjust = 0.5),
        legend.position = "top") +
  labs(title = "Mitochondrial volume distribution by group",
       x = "Volume (µm³)", y = "Relative frequency",
       fill = "Group", color = "Group")

print(p_vol_distribution)
ggsave(file.path(output_path, "volume_distribution_groups.png"),
       p_vol_distribution, width = 8, height = 5, dpi = 300)


################################################################################
## 4. CALCULATING MORPHOLOGICAL SUBTYPE CLUSTERS                                                    
##                                                                            
## Mitochondria are classified into morphological subtypes using k-means      
## clustering on scaled morphometric features. The optimal number of          
## clusters is assessed with the elbow method on a random subsample           
## (n = 10,000). A Random Forest classifier is then trained on the k-means   
## labels to produce a decision boundary and to rank feature importance       


feature_cols <- c(
  "Volume", "SA_Mesh", "Sph_Mesh", "RelativeVolume",
  "SA_to_Vol", "Equiv_diameter", "Compactness",
  "Vol_discrepancy", "SA_discrepancy", "Sph_discrepancy"
)

features_scaled_filt <- morpho_df_filtered %>%
  select(all_of(feature_cols)) %>%
  scale() %>%
  as.data.frame()

cluster_seed <- 123  # for reproducibility

# Elbow plot to guide choice of k
set.seed(cluster_seed)
sample_idx <- sample(nrow(features_scaled_filt), min(10000, nrow(features_scaled_filt)))
sample_df  <- features_scaled_filt[sample_idx, ]

wss <- sapply(1:8, function(k) kmeans(sample_df, centers = k, nstart = 10)$tot.withinss)

p_elbow <- data.frame(k = 1:8, WSS = wss) %>%
  ggplot(aes(x = k, y = WSS)) +
  geom_line(color = "darkgray", linewidth = 1.2) +
  geom_point(size = 4, color = "black") +
  theme_classic() +
  theme(axis.title = element_text(size = 14), axis.text = element_text(size = 12),
        plot.title = element_text(size = 15, face = "bold", hjust = 0.5)) +
  labs(title = "Elbow method — optimal k",
       x = "Number of clusters (k)", y = "Total within-cluster SS")

print(p_elbow)
optimal_k    <- 4    # number of morphological clusters (inspect elbow plot first and choose the best value for your analyses)

ggsave(file.path(output_path, "elbow_plot.png"), p_elbow, width = 6, height = 4, dpi = 300)

# k-means clustering on full filtered dataset
set.seed(cluster_seed)
kmeans_model <- kmeans(features_scaled_filt, centers = optimal_k,
                       nstart = 20, iter.max = 50)

morpho_df_filtered$Cluster <- factor(kmeans_model$cluster)

cluster_summary <- morpho_df_filtered %>%
  group_by(Cluster) %>%
  summarise(
    n               = n(),
    Mean_Volume     = mean(Volume),
    Mean_SA         = mean(SA_Mesh),
    Mean_Sphericity = mean(Sph_Mesh),
    Mean_SA_to_Vol  = mean(SA_to_Vol),
    Mean_Diameter   = mean(Equiv_diameter),
    .groups = "drop"
  ) %>%
  arrange(Mean_Volume)

cat("\n── Cluster summary (sorted by mean volume) ──\n")
print(cluster_summary)
cat("\nInspect cluster_summary, update subtype_recode in the configuration\n",
    "section, then re-run from here.\n\n")

# Assign biological subtype labels
# After running the clustering section and inspecting cluster_summary, assign
# a biological label to each cluster number. Labels must match subtype_levels.
subtype_recode <- c(
  "1" = "Fragmented",
  "4" = "Intermediate",
  "3" = "Elongated",
  "2" = "Networked") ## EDIT cluster names based on YOUR cluster_summary output
subtype_levels <- c("Fragmented", "Intermediate", "Elongated", "Networked")

morpho_df_filtered$Subtype <- recode(morpho_df_filtered$Cluster, subtype_recode)
morpho_df_filtered$Subtype <- factor(morpho_df_filtered$Subtype, levels = subtype_levels)

cat("── Subtype counts after labelling ──\n")

# UMAP visualization (subsampled for speed)
set.seed(42)
viz_idx     <- sample(nrow(features_scaled_filt), min(20000, nrow(features_scaled_filt)))
umap_result <- umap(features_scaled_filt[viz_idx, ])

umap_df <- data.frame(
  UMAP1   = umap_result$layout[, 1],
  UMAP2   = umap_result$layout[, 2],
  Subtype = morpho_df_filtered$Subtype[viz_idx],
  Group   = morpho_df_filtered$Group[viz_idx]
)

p_umap_subtype <- ggplot(umap_df, aes(x = UMAP1, y = UMAP2, color = Subtype)) +
  geom_point(alpha = 0.3, size = 0.6) +
  scale_color_brewer(palette = "Set2", na.value = "grey50") +
  theme_classic() +
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1))) +
  labs(title = "Mitochondrial subtypes — UMAP", color = "Subtype")

p_umap_group <- ggplot(umap_df, aes(x = UMAP1, y = UMAP2, color = Group)) +
  geom_point(alpha = 0.3, size = 0.6) +
  scale_color_manual(values = group_colors, na.value = "grey50") +
  theme_classic() +
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1))) +
  labs(title = "Experimental groups — UMAP", color = "Group")

print(p_umap_subtype | p_umap_group)
ggsave(file.path(output_path, "umap_visualization.png"),
       p_umap_subtype | p_umap_group, width = 12, height = 5, dpi = 300)

# Random Forest classifier
set.seed(42)
train_idx <- createDataPartition(morpho_df_filtered$Subtype, p = 0.8, list = FALSE)

X_train <- features_scaled_filt[ train_idx, ]
X_test  <- features_scaled_filt[-train_idx, ]
y_train <- morpho_df_filtered$Subtype[ train_idx]
y_test  <- morpho_df_filtered$Subtype[-train_idx]

rf_model <- randomForest(x = X_train, y = y_train, ntree = 300, importance = TRUE)

cat("\n── Random Forest performance on held-out test set ──\n")
pred <- predict(rf_model, X_test)
print(confusionMatrix(pred, y_test))

importance_df <- data.frame(
  Feature    = rownames(importance(rf_model)),
  Importance = importance(rf_model)[, "MeanDecreaseGini"]
) %>% arrange(Importance)

p_importance <- ggplot(importance_df,
                       aes(x = reorder(Feature, Importance), y = Importance)) +
  geom_col(fill = "#4E9BB9", alpha = 0.8) +
  coord_flip() +
  theme_classic() +
  labs(title = "Feature importance — Random Forest",
       x = "", y = "Mean Decrease Gini")

print(p_importance)
ggsave(file.path(output_path, "feature_importance.png"),
       p_importance, width = 8, height = 5, dpi = 300)

# Apply RF classifier to full filtered dataset
morpho_df_filtered$Subtype <- predict(rf_model, features_scaled_filt)
morpho_df_filtered$Subtype <- factor(morpho_df_filtered$Subtype, levels = subtype_levels)


################################################################################
## 5. CALCULATING MORPHOLOGICAL SUBTYPE PROPORTIONS              

# Per-sample subtype proportions
subtype_props <- morpho_df_filtered %>%
  group_by(source_file, Group, Subtype) %>%
  summarise(n = n(), .groups = "drop") %>%
  group_by(source_file) %>%
  mutate(Proportion = n / sum(n)) %>%
  ungroup() %>%
  complete(source_file, Subtype, fill = list(n = 0, Proportion = 0)) %>%
  select(-Group) %>%
  left_join(
    morpho_df_filtered %>% select(source_file, Group) %>% distinct(),
    by = "source_file"
  ) %>%
  mutate(Subtype = factor(Subtype, levels = subtype_levels))

# Stacked bar: mean proportions per group
p_stacked <- subtype_props %>%
  group_by(Group, Subtype) %>%
  summarise(Mean_proportion = mean(Proportion), .groups = "drop") %>%
  ggplot(aes(x = Group, y = Mean_proportion, fill = Subtype)) +
  geom_col(position = "stack", alpha = 0.85) +
  scale_fill_brewer(palette = "Set2") +
  scale_y_continuous(labels = percent_format()) +
  theme_classic() +
  labs(title = "Mitochondrial subtype composition",
       y = "Mean proportion", x = "", fill = "Subtype")

# Boxplot: per-sample proportions by subtype
p_box <- ggplot(subtype_props, aes(x = Subtype, y = Proportion, fill = Group)) +
  geom_boxplot(position = position_dodge(width = 0.7), width = 0.5,
               outlier.shape = NA, alpha = 0.4, color = "grey30", linewidth = 0.8) +
  geom_jitter(aes(color = Group),
              position = position_jitterdodge(jitter.width = 0.15, dodge.width = 0.7),
              size = 2, alpha = 0.8) +
  stat_compare_means(aes(group = Group), method = "wilcox.test",
                     label = "p.signif", size = 6) +
  scale_y_continuous(labels = percent_format()) +
  scale_fill_manual(values  = group_colors) +
  scale_color_manual(values = dot_colors) +
  theme_classic() +
  theme(legend.position = "bottom",
        axis.title = element_text(size = 14), axis.text = element_text(size = 12),
        plot.title = element_text(size = 15, hjust = 0)) +
  labs(title = "Subtype proportions per sample", x = "Subtype", y = "Proportion")

print(p_stacked)
print(p_box)
ggsave(file.path(output_path, "subtype_proportions.png"),
       p_stacked / p_box, width = 12, height = 8, dpi = 300)

# Wilcoxon tests with Benjamini-Hochberg correction
group_levels <- unique(metadata$Group)

subtype_stats <- subtype_props %>%
  group_by(Subtype) %>%
  summarise(
    p_value = tryCatch(
      wilcox.test(
        Proportion[Group == group_levels[1]],
        Proportion[Group == group_levels[2]]
      )$p.value,
      error = function(e) NA_real_
    ),
    .groups = "drop"
  ) %>%
  mutate(
    p_adjusted  = p.adjust(p_value, method = "BH"),
    Significant = p_adjusted < 0.05
  )

cat("\n── Subtype proportion statistics (Wilcoxon, BH-corrected) ──\n")
print(subtype_stats)


################################################################################
## 6. EXTRACT AND SAVE DATA                                                             

write.csv(morpho_df_filtered, file.path(output_path, "mitochondria_with_subtypes.csv"),  row.names = FALSE)
write.csv(subtype_props,      file.path(output_path, "subtype_proportions.csv"),         row.names = FALSE)
write.csv(subtype_stats,      file.path(output_path, "subtype_statistics.csv"),          row.names = FALSE)
write.csv(cluster_summary,    file.path(output_path, "cluster_summary.csv"),             row.names = FALSE)
write.csv(retention_summary,  file.path(output_path, "filtering_retention_summary.csv"), row.names = FALSE)

message("\nDone! All outputs saved to: ", output_path)