#!/usr/bin/env Rscript
library(dplyr)
library(tidyr)
library(ggplot2)
library(tools)

source("/home/bayegy/pipelines/metagenome/transcriptome/scripts/rUtils.R")

norm_abc <- function(df) {
  df <- df[, sapply(df, is.numeric), drop = FALSE]
  rownms <- gsub(
    "[^\\w\\- ]", # 与/home/bayegy/pipelines/metagenome/MetaGenome/scripts/lr_rpc.py对应
    "",
    rownames(df),
    perl = TRUE
  )
  not_dup <- !duplicated(rownms)
  df <- df[not_dup, , drop = FALSE]
  rownames(df) <- rownms[not_dup]
  return(
    data.frame(
      t(df),
      check.names = FALSE,
      stringsAsFactors = FALSE
    )
  )
}

lr <- function(xPath,
               xColumn,
               yColumn,
               yPath = "None",
               map = "None",
               category = "None",
               colors = "black,red,green,blue",
               out_dir = "./",
               swap_xy = FALSE) {
  # 读取数据
  colors <- unlist(strsplit(colors, ","))
  sum_features <- "sum_features"
  sum_features_x <- (xColumn == sum_features)
  sum_features_y <- (yColumn == sum_features)

  x_data <- norm_abc(
    read.input(xPath, header = TRUE, row.names = 1)
  )
  merged_data <- data.frame(
    Sample = rownames(x_data),
    check.names = FALSE,
    stringsAsFactors = FALSE
  )

  if (sum_features_x) {
    merged_data$X <- rowSums(x_data)
  } else {
    # print(xColumn)
    # print(head(colnames(x_data)))
    merged_data$X <- x_data[[xColumn]]
  }


  # Y <- NULL
  if (yPath != "None") {
    y_data <- norm_abc(
      read.input(yPath, header = TRUE, row.names = 1)
    )
    y_data <- y_data[merged_data$Sample, , drop = FALSE]
    if (sum_features_y) {
      merged_data$Y <- rowSums(y_data)
    } else if (
      yColumn %in% colnames(y_data)
    ) {
      merged_data$Y <- y_data[[yColumn]]
    }
  }

  if (map != "None") {
    map_data <- read.input(map, header = TRUE, row.names = 1)
    map_data <- map_data[merged_data$Sample, , drop = FALSE]
    if (
      (
        !"Y" %in% colnames(merged_data)
      ) && (
        yColumn %in% colnames(map_data)
      )
    ) {
      merged_data$Y <- map_data[[yColumn]]
    }
    if (
      category %in% colnames(map_data)
    ) {
      merged_data$category <- map_data[[category]]
    }
  }

  if (!"category" %in% colnames(merged_data)) {
    merged_data$category <- "AllSamples"
  }

  if (!"Y" %in% colnames(merged_data)) {
    if (yColumn %in% colnames(x_data)) {
      merged_data$Y <- x_data[[yColumn]]
    } else {
      stop("%s not found!", yColumn)
    }
  }

  merged_data <- merged_data[, c("X", "Y", "category")]
  merged_data <- na.omit(merged_data)

  merged_data <- merged_data[order(merged_data$category), , drop = FALSE]

  groups <- sort(unique(merged_data$category))

  # 按照分组进行线性回归分析
  my_lm <- function(group) {
    # print(.data)
    # browser()

    sum_res <- summary(
      stats::lm(
        Y ~ X,
        data = merged_data[
          merged_data$category == group, ,
          drop = FALSE
        ]
      )
    )
    pvalue <- round(sum_res$coefficients[2, 4], 5)
    label <- paste0(
      group,
      " Y=",
      round(sum_res$coefficients[2, 1], 5),
      "*X+",
      round(sum_res$coefficients[1, 1], 5),
      " R2=",
      round(sum_res$r.squared, 5),
      ifelse(pvalue < 1e-5, " P<1e-5", paste0(" P=", pvalue))
    )
    return(label)
  }

  # linear_regression <- merged_data %>%
  #   group_by(category) %>%
  #   summarize(
  #     label = my_lm(formula = Y ~ X)
  #   )

  # 保存线性回归结果表格
  # result_table <- paste0(out_dir, "/result_table.txt")
  # write.table(linear_regression, result_table, sep = "\t", row.names = FALSE)

  labels <- sapply(groups, my_lm)

  path2label <- function(path) {
    path <- tools::file_path_sans_ext(path)
    nameList <- unlist(strsplit(path, "/"))
    fileName <- nameList[length(nameList)]
    fileName <- gsub("abundance_unstratified", "", fileName, ignore.case = TRUE)
    fileName <- gsub("All", "", fileName, ignore.case = TRUE)
    # fileName <- gsub("detail", "", fileName)
    fileName <- gsub("[^0-9A-Za-z]+", " ", fileName)
    fileName <- trimws(fileName)
    fileName <- paste0("Sum of ", fileName)
    return(fileName)
  }

  xLabel <- ifelse(xColumn == "sum_features" && xPath != "None", path2label(xPath), xColumn)
  yLabel <- ifelse(yColumn == "sum_features" && yPath != "None", path2label(yPath), yColumn)

  # 绘制拟合曲线图
  plot <- ggplot(data = merged_data, aes(x = X, y = Y, color = category)) +
    scale_color_manual(
      values = colors,
      labels = labels
    ) +
    geom_point() +
    geom_smooth(method = "lm", se = FALSE) +
    labs(x = xLabel, y = yLabel) +
    theme_minimal() +
    theme(
      legend.position = "top",
      legend.direction = "vertical",
      legend.box = "vertical",
      legend.title = element_blank()
    )
  # theme(legend.position = "top")
  h <- 7 + length(groups) * 0.4

  ggsave(
    paste0(out_dir, "/LM.svg"),
    plot,
    width = 7,
    height = ifelse(h > 50, 50, h)
  )
}

fire(lr)