library(ggplot2)
library(dplyr)
library(gridExtra)
library(plotly)
library(htmlwidgets)

# ── logging setup ─────────────────────────────────────────────────────────────
log_file <- paste0("PCA_", format(Sys.time(), "%Y%m%d_%H%M%S"), ".log")
log_con  <- file(log_file, open="wt")

log_msg <- function(...) {
  msg <- paste0("[", format(Sys.time(), "%Y-%m-%d %H:%M:%S"), "] ", ...)
  message(msg)          # to console
  cat(msg, "\n", file=log_con, append=TRUE)  # to log file
}

log_msg("Script started")
log_msg("Log file: ", log_file)
# ──────────────────────────────────────────────────────────────────────────────

#ATAC-seq
log_msg("Loading ATAC-seq data...")
ATAC_bl_adt    <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/ATAC/ocr_blood_adt", header=T, row.names=1)
ATAC_br_adt    <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/ATAC/ocr_brain_adt", header=T, row.names=1)
ATAC_liver_adt <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/ATAC/ocr_liver_adt", header=T, row.names=1)
log_msg("ATAC loaded  -- ocr_blood_adt: ", ncol(ATAC_bl_adt),    " samples, ", nrow(ATAC_bl_adt),    " peaks")
log_msg("ATAC loaded  -- ocr_brain_adt: ", ncol(ATAC_br_adt),    " samples, ", nrow(ATAC_br_adt),    " peaks")
log_msg("ATAC loaded  -- ocr_liver_adt: ", ncol(ATAC_liver_adt), " samples, ", nrow(ATAC_liver_adt), " peaks")

# keep only BA lab samples
ATAC_bl_adt    <- ATAC_bl_adt[,    grepl("^BA_", colnames(ATAC_bl_adt)),    drop=FALSE]
ATAC_br_adt    <- ATAC_br_adt[,    grepl("^BA_", colnames(ATAC_br_adt)),    drop=FALSE]
ATAC_liver_adt <- ATAC_liver_adt[, grepl("^BA_", colnames(ATAC_liver_adt)), drop=FALSE]
log_msg("ATAC after BA filter -- ocr_blood_adt: ", ncol(ATAC_bl_adt),    " samples")
log_msg("ATAC after BA filter -- ocr_brain_adt: ", ncol(ATAC_br_adt),    " samples")
log_msg("ATAC after BA filter -- ocr_liver_adt: ", ncol(ATAC_liver_adt), " samples")

#RNA-seq
log_msg("Loading RNA-seq data...")
RNA_bl_adt    <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/RNA/rc_blood_adt", header=T, row.names=1)
RNA_br_adt    <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/RNA/rc_brain_adt", header=T, row.names=1)
RNA_liver_adt <- read.table("/BRC/shuhua/target/BPA_UPenn/raw_data/RNA/rc_liver_adt", header=T, row.names=1)
log_msg("RNA loaded   -- rc_blood_adt: ", ncol(RNA_bl_adt),    " samples, ", nrow(RNA_bl_adt),    " genes")
log_msg("RNA loaded   -- rc_brain_adt: ", ncol(RNA_br_adt),    " samples, ", nrow(RNA_br_adt),    " genes")
log_msg("RNA loaded   -- rc_liver_adt: ", ncol(RNA_liver_adt), " samples, ", nrow(RNA_liver_adt), " genes")

# keep only BA lab samples
RNA_bl_adt    <- RNA_bl_adt[,    grepl("^BA_", colnames(RNA_bl_adt)),    drop=FALSE]
RNA_br_adt    <- RNA_br_adt[,    grepl("^BA_", colnames(RNA_br_adt)),    drop=FALSE]
RNA_liver_adt <- RNA_liver_adt[, grepl("^BA_", colnames(RNA_liver_adt)), drop=FALSE]
log_msg("RNA after BA filter  -- rc_blood_adt: ", ncol(RNA_bl_adt),    " samples")
log_msg("RNA after BA filter  -- rc_brain_adt: ", ncol(RNA_br_adt),    " samples")
log_msg("RNA after BA filter  -- rc_liver_adt: ", ncol(RNA_liver_adt), " samples")


#####################
# color by exposure (3 groups), shape by sex (circle=F, diamond=M)
expo_colors <- c("Ctrl"    = "#378ADD",
                 "BPA10mg" = "#D50032",
                 "BPA10ug" = "#FD7901")
sex_shapes_2d <- c("F" = 16, "M" = 18)   # ggplot: 16=circle, 18=diamond

group_colors <- c(
  "Ctrl_F"    = "#378ADD", "Ctrl_M"    = "#378ADD",
  "BPA10mg_F" = "#D50032", "BPA10mg_M" = "#D50032",
  "BPA10ug_F" = "#FD7901", "BPA10ug_M" = "#FD7901"
)
group_symbols <- c(
  "Ctrl_F"    = "circle",  "Ctrl_M"    = "diamond",
  "BPA10mg_F" = "circle",  "BPA10mg_M" = "diamond",
  "BPA10ug_F" = "circle",  "BPA10ug_M" = "diamond"
)

get_plot <- function(combined_data, my_tissue){

  log_msg("[", my_tissue, "] Starting PCA")

  filter <- apply(combined_data, 1, function(x) length(x[x>=5]) >= 5)
  combined_data1 <- combined_data[filter, ]
  log_msg("[", my_tissue, "] Features after count filter (>=5 in >=5 samples): ",
          nrow(combined_data1), " / ", nrow(combined_data))

  d <- apply(combined_data1+1, 2, function(x) x/sum(x)*10^6)
  m <- t(log2(d))

  pca <- prcomp(m, scale.=T)
  pca_plot <- as.data.frame(pca$x[, 1:3])

  sex <- rep("M", nrow(m))
  sex[grepl("_F_", rownames(m))] <- "F"
  pca_plot$sex <- sex

  pca_plot$expo <- as.character(data.frame(strsplit(colnames(combined_data1), "_"))[3,])

  var_perc <- round(pca$sdev^2 / sum(pca$sdev^2) * 100, digits=2)
  pca_plot$pc1_perc <- var_perc[1]
  pca_plot$pc2_perc <- var_perc[2]
  pca_plot$pc3_perc <- var_perc[3]
  log_msg("[", my_tissue, "] Variance explained -- PC1: ", var_perc[1],
          "%, PC2: ", var_perc[2], "%, PC3: ", var_perc[3], "%")

  #2D PCA
  p1 <- ggplot(pca_plot, aes(x=PC1, y=PC2, shape=sex, col=expo, label=rownames(pca_plot))) +
    geom_point(size=4) +
    scale_color_manual(values=expo_colors) +
    scale_shape_manual(values=sex_shapes_2d) +
    xlab(paste0("PC1 (", unique(pca_plot$pc1_perc), "%)")) +
    ylab(paste0("PC2 (", unique(pca_plot$pc2_perc), "%)")) +
    theme(title=element_text(size=10, face="bold")) +
    theme(legend.text=element_text(size=10)) +
    ggtitle(paste0("Controls across stages in ", my_tissue, " PC1 vs PC2 n=", nrow(combined_data1)))

  p2 <- ggplot(pca_plot, aes(x=PC1, y=PC3, shape=sex, col=expo, label=rownames(pca_plot))) +
    geom_point(size=4) +
    scale_color_manual(values=expo_colors) +
    scale_shape_manual(values=sex_shapes_2d) +
    xlab(paste0("PC1 (", unique(pca_plot$pc1_perc), "%)")) +
    ylab(paste0("PC3 (", unique(pca_plot$pc3_perc), "%)")) +
    theme(title=element_text(size=10, face="bold")) +
    theme(legend.text=element_text(size=10)) +
    ggtitle(paste0("Samples across stages in ", my_tissue, " PC1 vs PC3"))

  p3 <- ggplot(pca_plot, aes(x=PC2, y=PC3, shape=sex, col=expo, label=rownames(pca_plot))) +
    geom_point(size=4) +
    scale_color_manual(values=expo_colors) +
    scale_shape_manual(values=sex_shapes_2d) +
    xlab(paste0("PC2 (", unique(pca_plot$pc2_perc), "%)")) +
    ylab(paste0("PC3 (", unique(pca_plot$pc3_perc), "%)")) +
    theme(title=element_text(size=10, face="bold")) +
    theme(legend.text=element_text(size=10)) +
    ggtitle(paste0("Samples across stages in ", my_tissue, " PC2 vs PC3"))

  out_dir <- "/home/yan/TaRGET_II/BPA_brian/PCA"
  dir.create(out_dir, recursive=TRUE, showWarnings=FALSE)

  pdf_out <- file.path(out_dir, paste0("pca_UPenn_", my_tissue, ".pdf"))
  pdf(pdf_out, width=15, height=14)
  grid.arrange(p1, p2, p3, nrow=2)
  dev.off()
  log_msg("[", my_tissue, "] 2D PCA saved -> ", pdf_out)

  ####### 3D PCA
  pca_plot$sample_id <- rownames(pca_plot)
  pca_plot$group     <- paste(pca_plot$expo, pca_plot$sex, sep="_")

  fig <- plot_ly()
  for (grp in names(group_colors)) {
    df_sub <- pca_plot[pca_plot$group == grp, ]
    if (nrow(df_sub) == 0) next
    fig <- fig %>%
      add_trace(
        data          = df_sub,
        x = ~PC1, y = ~PC2, z = ~PC3,
        type          = "scatter3d",
        mode          = "markers",
        name          = grp,
        text          = ~sample_id,
        hovertemplate = "<b>%{text}</b><br>PC1: %{x:.2f}<br>PC2: %{y:.2f}<br>PC3: %{z:.2f}<extra></extra>",
        marker        = list(
          size    = 10,
          color   = group_colors[grp],
          symbol  = group_symbols[grp],
          opacity = 0.88,
          line    = list(color="rgba(255,255,255,0.5)", width=0.5)
        )
      )
  }

  fig <- fig %>%
    layout(
      title = paste0("3D PCA - ", my_tissue),
      scene = list(
        xaxis = list(title=paste0("PC1 (", unique(pca_plot$pc1_perc), "%)")),
        yaxis = list(title=paste0("PC2 (", unique(pca_plot$pc2_perc), "%)")),
        zaxis = list(title=paste0("PC3 (", unique(pca_plot$pc3_perc), "%)"))
      )
    )

  html_out <- file.path(out_dir, paste0("pca_", my_tissue, ".html"))
  htmlwidgets::saveWidget(as_widget(fig), selfcontained=TRUE, html_out)
  log_msg("[", my_tissue, "] 3D PCA saved -> ", html_out)
}

get_plot(ATAC_bl_adt,    "ATAC_bl_adt")
get_plot(ATAC_br_adt,    "ATAC_br_adt")
get_plot(ATAC_liver_adt, "ATAC_liver_adt")
get_plot(RNA_bl_adt,     "RNA_bl_adt")
get_plot(RNA_br_adt,     "RNA_br_adt")
get_plot(RNA_liver_adt,  "RNA_liver_adt")

log_msg("All done")
close(log_con)