#!/usr/bin/env r
#**************************************************
## Plot Heatmap showing the effect of the compounds
## in transcriptomics, metabolomics, and lipidomics
## modalities
#**************************************************

#-------------------------
## Load required libraries
#-------------------------
library(tidyverse)
library(vroom)

#------------------------
## Plotting Variables
#------------------------
source("~/bin/theme_publication.R")
orgs <- c("CElegans", "Daphnia", "Drosophila", "Danio", "Xenopus", "HepG2")

#-------------------
## Working directory
#-------------------
setwd("/users/rg/vntasis/PTox/production_phase2_data/reports/")

#-----------------------
## Load and process data
#-----------------------
compound_names <- vroom('compounds.tsv')

# Annotated metabolites
annotations_dir <- "/users/project/precisiontox_008279_no_backup/production_metabolomics_data/annotations/"
annotations <-
  list.files(annotations_dir) %>%
  map(\(x) vroom(
    paste0(annotations_dir, x),
    col_types = "ccccccccclciccc",
    progress = FALSE
  )) %>%
  list_rbind()

# Results HILIC metabolomics
metabolomics_dir <- "/users/project/precisiontox_008279_no_backup/production_metabolomics_data/HILIC/"
metabolomics <-
  list.files(metabolomics_dir) %>%
  map(\(x) vroom(
    paste0(metabolomics_dir, x),
    col_types = "cccccccdddd",
    progress = FALSE
  )) %>%
  list_rbind() %>%
  filter(compound_id != "000") %>%
  left_join(
    annotations %>%
      filter(assay_type == "HILIC") %>%
      select(metabolite_id, biosystem_id, exposure_batch_id) %>%
      mutate(annotated = TRUE) %>%
      distinct()
  ) %>%
  group_by(biosystem_id, exposure_batch_id, timepoint_level, compound_id, dose_code) %>%
  summarise(
    N_signif = sum(abs(ratio) >= log2(1.5) & padj <= 0.05, na.rm = TRUE),
    N_signif_an = sum(annotated & abs(ratio) >= log2(1.5) & padj <= 0.05, na.rm = TRUE)
  ) %>%
  ungroup()

# Results lipidomics
lipidomics_dir <- "/users/project/precisiontox_008279_no_backup/production_metabolomics_data/LIPIDS/"
lipidomics <-
  list.files(lipidomics_dir) %>%
  map(\(x) vroom(
    paste0(lipidomics_dir, x),
    col_types = "cccccccdddd",
    progress = FALSE
  )) %>%
  list_rbind() %>%
  filter(compound_id != "000") %>%
  left_join(
    annotations %>%
      filter(assay_type == "LIPIDS") %>%
      select(metabolite_id, biosystem_id, exposure_batch_id) %>%
      mutate(annotated = TRUE) %>%
      distinct()
  ) %>%
  group_by(biosystem_id, exposure_batch_id, timepoint_level, compound_id, dose_code) %>%
  summarise(
    N_signif = sum(abs(ratio) >= log2(1.5) & padj <= 0.05, na.rm = TRUE),
    N_signif_an = sum(annotated & abs(ratio) >= log2(1.5) & padj <= 0.05, na.rm = TRUE)
  ) %>%
  ungroup()

# Results transcriptomics
diff_ab_wald <-
  vroom('differential_abundance_summary_Wald.tsv') %>%
  select(-N_Terms) %>%
  pivot_wider(names_from = "compound", values_from = "N_DEGs") %>%
  pivot_longer(
    -model,
    names_to = "compound",
    values_to = "N_DEGs"
  ) %>%
  separate_wider_delim(compound, "_", names = c("compound_id", "timepoint_level"))

# Merge in one table
to_plot <-
  metabolomics %>%
  rename(N_met = N_signif, N_met_an = N_signif_an) %>%
  full_join(
    lipidomics %>%
      rename(N_lip = N_signif, N_lip_an = N_signif_an)
  ) %>%
  mutate(
    model = case_when(
      biosystem_id == "F" ~ "Drosophila",
      biosystem_id == "R" ~ "Danio",
      biosystem_id == "E" ~ "CElegans",
      biosystem_id == "D" ~ "Daphnia",
      biosystem_id == "X" ~ "Xenopus",
      biosystem_id == "G" ~ "HepG2",
    )
  ) %>%
  inner_join(diff_ab_wald) %>%
  mutate(
    Compound_code = as.double(compound_id),
    timepoint = factor(timepoint_level, levels = c("A", "B"), labels = c("6 hrs", "24 hrs")),
    model = factor(model, levels = orgs)
  ) %>%
  left_join(compound_names) %>%
  select(
    Compound_name, model, timepoint,
    exposure_batch_id, dose_code,
    N_DEGs, N_met, N_met_an, N_lip, N_lip_an
  ) %>%
  pivot_longer(
    N_DEGs:N_lip_an,
    names_to = "modality",
    values_to = "N_signif"
  ) %>%
  mutate(
    modality = factor(
      modality,
      levels = c("N_DEGs", "N_met", "N_met_an", "N_lip", "N_lip_an"),
      labels = c("Transcriptomics", "Metabolomics", "Metabolomics_Ant", "Lipidomics", "Lipidomics_Ant")
    )
  )

# Introduce NA lines where data is not complete
to_plot <-
  to_plot %>%
  select(Compound_name) %>%
  distinct() %>%
  mutate(include = 1) %>%
  left_join(
    to_plot %>%
      select(model) %>%
      distinct() %>%
      mutate(include = 1),
    relationship = "many-to-many"
  ) %>%
  left_join(
    to_plot %>%
      select(modality) %>%
      distinct() %>%
      mutate(include = 1),
    relationship = "many-to-many"
  ) %>%
  left_join(
    to_plot %>%
      select(timepoint) %>%
      distinct() %>%
      mutate(include = 1),
    relationship = "many-to-many"
  ) %>%
  select(-include) %>%
  left_join(to_plot)

# Group according to exposure batch and normalize signal
# row min-max norm
# col min-max norm
# capping at 2000
to_plot <-
  to_plot %>%
  left_join(
    to_plot %>%
      select(Compound_name, exposure_batch_id) %>%
      distinct() %>%
      group_by(Compound_name) %>%
      summarise(exposure_batch =
        exposure_batch_id %>%
        na.omit() %>%
        paste0(collapse = ",")
      )
  ) %>%
  group_by(Compound_name, modality) %>%
  mutate(
    N_signif_scaled_row =
      (N_signif - min(N_signif, na.rm = TRUE)) / (max(N_signif, na.rm = TRUE) - min(N_signif, na.rm = TRUE))
  ) %>%
  group_by(model, timepoint, modality) %>%
  mutate(
    N_signif_scaled_col =
      (N_signif - min(N_signif, na.rm = TRUE)) / (max(N_signif, na.rm = TRUE) - min(N_signif, na.rm = TRUE))
  ) %>%
  rowwise() %>%
  mutate(
    N_signif_capped = min(N_signif, 2000)
  ) %>%
  ungroup()



#---------------
## Generate plot
#---------------
pdf(
  "metabolomics_transcriptomics_heatmap.pdf",
  height = 9,
  width = 12
)

to_plot %>%
  filter(!(modality %in% c("Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title = "Raw value of Significant features"
  ) +
  scale_fill_distiller("# Significant", palette = "PuBu", direction = 1)

to_plot %>%
  filter(!(modality %in% c("Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif_capped)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title = "Raw value of Significant features capped at 2000"
  ) +
  scale_fill_distiller("# Significant", palette = "PuBu", direction = 1)

to_plot %>%
  filter(!(modality %in% c("Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif_scaled_row)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title = "Value of Significant features scaled per compound and modality (row)\nComparison across Species"
  ) +
  scale_fill_distiller("# Significant\n   scaled", palette = "PuBu", direction = 1)

to_plot %>%
  filter(!(modality %in% c("Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif_scaled_col)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title =
      "Value of Significant features scaled per Species, time, and modality (column)\nComparison across compounds"
  ) +
  scale_fill_distiller("# Significant\n   scaled", palette = "PuBu", direction = 1)

to_plot %>%
  filter((modality %in% c("Transcriptomics", "Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif_scaled_row)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title = "Value of Significant features scaled per compound and modality (row)\nfor annotated metabolome and lipidome compounds"
  ) +
  scale_fill_distiller("# Significant\n   scaled", palette = "PuBu", direction = 1)

to_plot %>%
  filter((modality %in% c("Transcriptomics", "Metabolomics_Ant", "Lipidomics_Ant"))) %>%
  ggplot(aes(model, Compound_name, fill = N_signif_scaled_col)) +
  geom_tile(colour = "grey") +
  facet_grid(exposure_batch ~ modality + timepoint, scales = "free_y") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.text.x = element_text(angle = 90, vjust = 1, hjust = 1),
    strip.text.y = element_blank()
  ) +
  labs(
    x = "",
    y = "Compound",
    title =
      "Value of Significant features scaled per Species, time, and modality (column)\nfor annotated metabolome and lipidome compounds"
  ) +
  scale_fill_distiller("# Significant\n   scaled", palette = "PuBu", direction = 1)


dev.off()