#' Draw marker-gene heatmap across annotated cell groups
#'
#' @description
#' `CellWindX_MarkerHeatmap()` visualizes marker gene expression statistics
#' across annotated cell groups using either `ComplexHeatmap` or `ggplot2`.
#'
#' @param marker_result A result object returned by `CellWindX_TopMarkersStats()`,
#'   or a data frame containing marker expression statistics.
#' @param value_col Character string. Value column used for heatmap visualization.
#'   One of `"avg_expr"`, `"pct_expr"`, or `"avg_expr_positive"`.
#' @param scale_method Character string. Scaling method. One of `"zscore"` or `"none"`.
#' @param plot_engine Character string. Plotting engine. One of `"complex"` or `"ggplot"`.
#' @param palette Character string. Built-in CellWindX palette.
#' @param cluster_rows Logical. Whether to cluster rows.
#' @param cluster_columns Logical. Whether to cluster columns.
#' @param show_row_names Logical. Whether to show row names.
#' @param show_column_names Logical. Whether to show column names.
#' @param column_names_rot Numeric. Rotation angle of column names.
#' @param round_cell Logical. Whether to draw rounded cells when using ComplexHeatmap.
#' @param cell_width_mm Numeric. Cell width in millimeters.
#' @param cell_height_mm Numeric. Cell height in millimeters.
#' @param row_fontsize Numeric. Row name font size.
#' @param column_fontsize Numeric. Column name font size.
#' @param legend_position Character string. Legend position.
#' @param heatmap_title Character string or NULL. Heatmap title.
#' @param zscore_clip Numeric. Z-score clipping threshold.
#' @param output_file Character string or NULL. Optional output file path.
#' @param output_width Numeric. Output width in inches.
#' @param output_height Numeric. Output height in inches.
#' @param dpi Numeric. Output resolution.
#' @param draw Logical. Whether to draw the heatmap immediately when using ComplexHeatmap.
#'
#' @return A list with class `"CellWindX_MarkerHeatmap"`.
#'
#' @export
CellWindX_MarkerHeatmap <- function(
    marker_result,
    value_col = c("avg_expr", "pct_expr", "avg_expr_positive"),
    scale_method = c("zscore", "none"),
    plot_engine = c("complex", "ggplot"),
    palette = c("shanshui", "chongqing_modern", "girlish"),
    cluster_rows = FALSE,
    cluster_columns = FALSE,
    show_row_names = TRUE,
    show_column_names = TRUE,
    column_names_rot = 45,
    round_cell = TRUE,
    cell_width_mm = 5.5,
    cell_height_mm = 5.5,
    row_fontsize = 10,
    column_fontsize = 9,
    legend_position = "right",
    heatmap_title = NULL,
    zscore_clip = 2,
    output_file = NULL,
    output_width = 9,
    output_height = 5,
    dpi = 300,
    draw = TRUE
) {

  # -----------------------------
  # 1. Package check
  # -----------------------------
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("Package 'ggplot2' is required.")
  }
  if (!requireNamespace("dplyr", quietly = TRUE)) {
    stop("Package 'dplyr' is required.")
  }
  if (!requireNamespace("tidyr", quietly = TRUE)) {
    stop("Package 'tidyr' is required.")
  }

  value_col <- match.arg(value_col)
  scale_method <- match.arg(scale_method)
  plot_engine <- match.arg(plot_engine)
  palette <- match.arg(palette)

  if (plot_engine == "complex") {
    if (!requireNamespace("ComplexHeatmap", quietly = TRUE)) {
      stop("Package 'ComplexHeatmap' is required for plot_engine = 'complex'.")
    }
    if (!requireNamespace("circlize", quietly = TRUE)) {
      stop("Package 'circlize' is required for plot_engine = 'complex'.")
    }
    if (!requireNamespace("grid", quietly = TRUE)) {
      stop("Package 'grid' is required for plot_engine = 'complex'.")
    }
  }

  # -----------------------------
  # 2. Built-in CellWindX palettes
  # -----------------------------
  CellWindX_palettes <- list(

    shanshui = list(
      discrete = c(
        "#8FB9A8", "#B7C8A3", "#A7C7D9", "#7FA6C7", "#9D9AC7",
        "#C6B7D9", "#D8C7A3", "#BFA58A", "#8A9A8F", "#6F8F9F",
        "#A6B1C2", "#C9D6DF", "#D6C6B9", "#E2D8C3", "#9E8FB2",
        "#738A99", "#B5A789", "#AFC8BA", "#D1A7A0", "#6E7F80"
      ),
      heatmap_low = "#6F8F9F",
      heatmap_mid = "#F7F4EA",
      heatmap_high = "#B58C9E"
    ),

    chongqing_modern = list(
      discrete = c(
        "#2F4858", "#33658A", "#4A90A4", "#86BBD8", "#758E4F",
        "#F6AE2D", "#F26419", "#D64550", "#8E5572", "#5D5179",
        "#3A506B", "#1C7293", "#00A6A6", "#7D8491", "#B8B8D1",
        "#E0A458", "#C8553D", "#6D597A", "#355070", "#2A2D34"
      ),
      heatmap_low = "#355070",
      heatmap_mid = "#F7F7F7",
      heatmap_high = "#F26419"
    ),

    girlish = list(
      discrete = c(
        "#FF9AA2", "#FFB7B2", "#FFDAC1", "#E2F0CB", "#B5EAD7",
        "#C7CEEA", "#F6C1D5", "#D5AAFF", "#A0E7E5", "#B4F8C8",
        "#FBE7C6", "#F7A8B8", "#F3C4FB", "#BFD7FF", "#FFD6E0",
        "#FFABAB", "#FFC3A0", "#D4F0F0", "#CCE2CB", "#E7C6FF"
      ),
      heatmap_low = "#A0C4FF",
      heatmap_mid = "#FFFFFF",
      heatmap_high = "#FF8FAB"
    )
  )

  pal_use <- CellWindX_palettes[[palette]]

  # -----------------------------
  # 3. Extract input data
  # -----------------------------
  if (is.list(marker_result) && "marker_expr_by_group" %in% names(marker_result)) {
    expr_df <- marker_result$marker_expr_by_group
  } else if (is.data.frame(marker_result)) {
    expr_df <- marker_result
  } else {
    stop(
      "marker_result must be either the output list from CellWindX_TopMarkersStats() ",
      "or a data.frame containing marker expression statistics."
    )
  }

  required_cols <- c(
    "source_cluster",
    "target_cluster",
    "gene",
    value_col
  )

  missing_cols <- setdiff(required_cols, colnames(expr_df))

  if (length(missing_cols) > 0) {
    stop(
      "The following required columns are missing: ",
      paste(missing_cols, collapse = ", ")
    )
  }

  expr_df <- as.data.frame(expr_df)
  expr_df$source_cluster <- as.character(expr_df$source_cluster)
  expr_df$target_cluster <- as.character(expr_df$target_cluster)
  expr_df$gene <- as.character(expr_df$gene)
  expr_df[[value_col]] <- as.numeric(expr_df[[value_col]])

  # -----------------------------
  # 4. Add marker rank if available
  # -----------------------------
  gene_rank_df <- NULL

  if (is.list(marker_result) && "top_markers" %in% names(marker_result)) {

    top_markers <- marker_result$top_markers

    if (all(c("cluster", "gene") %in% colnames(top_markers))) {

      top_markers$cluster <- as.character(top_markers$cluster)
      top_markers$gene <- as.character(top_markers$gene)

      if (!"rank_in_cluster" %in% colnames(top_markers)) {
        top_markers$rank_in_cluster <- ave(
          seq_len(nrow(top_markers)),
          top_markers$cluster,
          FUN = seq_along
        )
      }

      gene_rank_df <- top_markers[, c("cluster", "gene", "rank_in_cluster")]
      colnames(gene_rank_df) <- c("source_cluster", "gene", "rank_in_cluster")
    }
  }

  if (is.null(gene_rank_df)) {
    gene_rank_df <- unique(expr_df[, c("source_cluster", "gene")])
    gene_rank_df$rank_in_cluster <- ave(
      seq_len(nrow(gene_rank_df)),
      gene_rank_df$source_cluster,
      FUN = seq_along
    )
  }

  # -----------------------------
  # 5. Create unique marker column ID
  # -----------------------------
  expr_df$marker_id <- paste(expr_df$source_cluster, expr_df$gene, sep = "__")
  gene_rank_df$marker_id <- paste(gene_rank_df$source_cluster, gene_rank_df$gene, sep = "__")

  expr_df <- expr_df[
    !duplicated(expr_df[, c("target_cluster", "marker_id")]),
    ,
    drop = FALSE
  ]

  # -----------------------------
  # 6. Construct matrix
  # -----------------------------
  mat_df <- expr_df[, c("target_cluster", "marker_id", value_col)]

  mat_wide <- tidyr::pivot_wider(
    mat_df,
    names_from = "marker_id",
    values_from = value_col
  )

  mat_wide <- as.data.frame(mat_wide)
  rownames(mat_wide) <- mat_wide$target_cluster
  mat_wide$target_cluster <- NULL

  mat <- as.matrix(mat_wide)
  storage.mode(mat) <- "numeric"
  mat[is.na(mat)] <- 0

  # -----------------------------
  # 7. Order rows and columns
  # -----------------------------
  target_order <- unique(expr_df$target_cluster)
  target_order <- target_order[target_order %in% rownames(mat)]

  gene_rank_df <- gene_rank_df[
    gene_rank_df$marker_id %in% colnames(mat),
    ,
    drop = FALSE
  ]

  gene_rank_df <- gene_rank_df[
    order(gene_rank_df$source_cluster, gene_rank_df$rank_in_cluster),
    ,
    drop = FALSE
  ]

  marker_order <- unique(gene_rank_df$marker_id)
  marker_order <- marker_order[marker_order %in% colnames(mat)]

  mat <- mat[target_order, marker_order, drop = FALSE]

  gene_info <- gene_rank_df[
    match(colnames(mat), gene_rank_df$marker_id),
    ,
    drop = FALSE
  ]

  gene_labels <- gene_info$gene
  source_cluster_vec <- gene_info$source_cluster
  names(source_cluster_vec) <- colnames(mat)

  # -----------------------------
  # 8. Scaling
  # -----------------------------
  if (scale_method == "zscore") {

    mat_scaled <- scale(mat)
    mat_scaled[is.na(mat_scaled)] <- 0

    if (!is.null(zscore_clip)) {
      mat_scaled[mat_scaled > zscore_clip] <- zscore_clip
      mat_scaled[mat_scaled < -zscore_clip] <- -zscore_clip
    }

    plot_mat <- mat_scaled
    legend_name <- "Z-score"
    color_breaks <- c(-zscore_clip, 0, zscore_clip)

  } else {

    plot_mat <- mat

    if (value_col == "pct_expr") {
      legend_name <- "Percent"
    } else {
      legend_name <- value_col
    }

    color_breaks <- c(
      min(plot_mat, na.rm = TRUE),
      stats::median(plot_mat, na.rm = TRUE),
      max(plot_mat, na.rm = TRUE)
    )
  }

  if (length(unique(color_breaks)) < 3) {
    color_breaks <- c(
      min(plot_mat, na.rm = TRUE),
      mean(plot_mat, na.rm = TRUE),
      max(plot_mat, na.rm = TRUE)
    )
    color_breaks <- jitter(color_breaks, factor = 1e-8)
  }

  # -----------------------------
  # 9. Cluster annotation colors
  # -----------------------------
  source_levels <- unique(source_cluster_vec)

  if (length(source_levels) > length(pal_use$discrete)) {
    stop("The built-in palette supports up to 20 groups.")
  }

  source_colors <- pal_use$discrete[seq_along(source_levels)]
  names(source_colors) <- source_levels

  if (is.null(heatmap_title)) {
    heatmap_title <- paste0(
      "CellWindX marker heatmap | ",
      value_col,
      " | ",
      scale_method,
      " | ",
      palette
    )
  }

  # -----------------------------
  # 10. ComplexHeatmap engine
  # -----------------------------
  if (plot_engine == "complex") {

    col_fun <- circlize::colorRamp2(
      color_breaks,
      c(
        pal_use$heatmap_low,
        pal_use$heatmap_mid,
        pal_use$heatmap_high
      )
    )

    top_anno <- ComplexHeatmap::HeatmapAnnotation(
      MarkerSource = ComplexHeatmap::anno_simple(
        source_cluster_vec,
        col = source_colors,
        height = grid::unit(2.2, "mm")
      ),
      show_annotation_name = FALSE
    )

    if (round_cell) {

      ht <- ComplexHeatmap::Heatmap(
        plot_mat,
        name = legend_name,
        col = col_fun,
        cluster_rows = cluster_rows,
        cluster_columns = cluster_columns,
        show_row_names = show_row_names,
        show_column_names = show_column_names,
        column_labels = gene_labels,
        column_names_rot = column_names_rot,
        column_split = factor(source_cluster_vec, levels = source_levels),
        top_annotation = top_anno,
        rect_gp = grid::gpar(col = NA, fill = NA),
        cell_fun = function(j, i, x, y, w, h, fill) {
          v <- plot_mat[i, j]
          grid::grid.roundrect(
            x = x,
            y = y,
            width = w * 0.90,
            height = h * 0.90,
            r = grid::unit(0.35, "snpc"),
            gp = grid::gpar(
              fill = col_fun(v),
              col = NA
            )
          )
        },
        use_raster = FALSE,
        width = ncol(plot_mat) * grid::unit(cell_width_mm, "mm"),
        height = nrow(plot_mat) * grid::unit(cell_height_mm, "mm"),
        row_names_gp = grid::gpar(fontsize = row_fontsize),
        column_names_gp = grid::gpar(fontsize = column_fontsize),
        column_title_gp = grid::gpar(fontsize = 9, fontface = "bold"),
        heatmap_legend_param = list(
          title = legend_name,
          title_gp = grid::gpar(fontsize = 9, fontface = "bold"),
          labels_gp = grid::gpar(fontsize = 8)
        )
      )

    } else {

      ht <- ComplexHeatmap::Heatmap(
        plot_mat,
        name = legend_name,
        col = col_fun,
        cluster_rows = cluster_rows,
        cluster_columns = cluster_columns,
        show_row_names = show_row_names,
        show_column_names = show_column_names,
        column_labels = gene_labels,
        column_names_rot = column_names_rot,
        column_split = factor(source_cluster_vec, levels = source_levels),
        top_annotation = top_anno,
        rect_gp = grid::gpar(col = "white", lwd = 0.7),
        width = ncol(plot_mat) * grid::unit(cell_width_mm, "mm"),
        height = nrow(plot_mat) * grid::unit(cell_height_mm, "mm"),
        row_names_gp = grid::gpar(fontsize = row_fontsize),
        column_names_gp = grid::gpar(fontsize = column_fontsize),
        column_title_gp = grid::gpar(fontsize = 9, fontface = "bold"),
        heatmap_legend_param = list(
          title = legend_name,
          title_gp = grid::gpar(fontsize = 9, fontface = "bold"),
          labels_gp = grid::gpar(fontsize = 8)
        )
      )
    }

    if (draw) {
      ComplexHeatmap::draw(
        ht,
        heatmap_legend_side = legend_position,
        annotation_legend_side = legend_position
      )
    }

    if (!is.null(output_file)) {

      if (grepl("\\.pdf$", output_file, ignore.case = TRUE)) {
        grDevices::pdf(output_file, width = output_width, height = output_height)
        ComplexHeatmap::draw(
          ht,
          heatmap_legend_side = legend_position,
          annotation_legend_side = legend_position
        )
        grDevices::dev.off()
      } else if (grepl("\\.png$", output_file, ignore.case = TRUE)) {
        grDevices::png(
          output_file,
          width = output_width,
          height = output_height,
          units = "in",
          res = dpi
        )
        ComplexHeatmap::draw(
          ht,
          heatmap_legend_side = legend_position,
          annotation_legend_side = legend_position
        )
        grDevices::dev.off()
      } else {
        warning("Only .pdf and .png output_file are supported for plot_engine = 'complex'.")
      }
    }

    result <- list(
      plot = ht,
      matrix_raw = mat,
      matrix_plot = plot_mat,
      gene_info = gene_info,
      source_colors = source_colors,
      palette = palette,
      plot_engine = plot_engine,
      value_col = value_col,
      scale_method = scale_method
    )

    class(result) <- c("CellWindX_MarkerHeatmap", class(result))
    return(result)
  }

  # -----------------------------
  # 11. ggplot2 engine
  # -----------------------------
  if (plot_engine == "ggplot") {

    plot_df <- as.data.frame(plot_mat)
    plot_df$target_cluster <- rownames(plot_mat)

    plot_df <- tidyr::pivot_longer(
      plot_df,
      cols = -target_cluster,
      names_to = "marker_id",
      values_to = "value"
    )

    gene_info2 <- gene_info[, c("marker_id", "source_cluster", "gene")]

    plot_df <- dplyr::left_join(
      plot_df,
      gene_info2,
      by = "marker_id"
    )

    plot_df$target_cluster <- factor(
      plot_df$target_cluster,
      levels = rev(rownames(plot_mat))
    )

    plot_df$marker_id <- factor(
      plot_df$marker_id,
      levels = colnames(plot_mat)
    )

    plot_df$gene <- factor(
      plot_df$gene,
      levels = unique(gene_info$gene)
    )

    plot_df$source_cluster <- factor(
      plot_df$source_cluster,
      levels = source_levels
    )

    p <- ggplot2::ggplot(
      plot_df,
      ggplot2::aes(
        x = marker_id,
        y = target_cluster,
        fill = value
      )
    ) +
      ggplot2::geom_tile(
        color = "white",
        linewidth = 0.35
      ) +
      ggplot2::scale_x_discrete(
        labels = stats::setNames(gene_info$gene, gene_info$marker_id)
      ) +
      ggplot2::scale_fill_gradient2(
        low = pal_use$heatmap_low,
        mid = pal_use$heatmap_mid,
        high = pal_use$heatmap_high,
        midpoint = 0,
        name = legend_name
      ) +
      ggplot2::facet_grid(
        . ~ source_cluster,
        scales = "free_x",
        space = "free_x"
      ) +
      ggplot2::labs(
        x = NULL,
        y = NULL,
        title = heatmap_title
      ) +
      ggplot2::theme_minimal(base_size = 12) +
      ggplot2::theme(
        plot.title = ggplot2::element_text(
          hjust = 0.5,
          face = "bold",
          size = 14
        ),
        panel.grid = ggplot2::element_blank(),
        axis.text.x = ggplot2::element_text(
          size = column_fontsize,
          angle = column_names_rot,
          hjust = 1,
          vjust = 1
        ),
        axis.text.y = ggplot2::element_text(
          size = row_fontsize,
          color = "black"
        ),
        strip.text.x = ggplot2::element_text(
          size = 10,
          face = "bold"
        ),
        panel.spacing.x = grid::unit(0.15, "cm"),
        legend.position = ifelse(
          legend_position %in% c("right", "left", "top", "bottom", "none"),
          legend_position,
          "right"
        ),
        legend.title = ggplot2::element_text(
          size = 9,
          face = "bold"
        ),
        legend.text = ggplot2::element_text(size = 8),
        plot.margin = ggplot2::margin(8, 8, 8, 8)
      )

    if (!is.null(output_file)) {
      ggplot2::ggsave(
        filename = output_file,
        plot = p,
        width = output_width,
        height = output_height,
        dpi = dpi
      )
    }

    result <- list(
      plot = p,
      matrix_raw = mat,
      matrix_plot = plot_mat,
      plot_data = plot_df,
      gene_info = gene_info,
      source_colors = source_colors,
      palette = palette,
      plot_engine = plot_engine,
      value_col = value_col,
      scale_method = scale_method
    )

    class(result) <- c("CellWindX_MarkerHeatmap", class(result))
    return(result)
  }
}
