#' Detect the type of the input data for Upset plot
#'
#' @keywords internal
#' @param data A data frame or a list
#' @param group_by A character string specifying the column name of the data frame to group the data.
#' @param id_by A character string specifying the column name of the data frame to identify the instances.
#'   Required when `group_by` is a single column and data is a data frame.
#' @return A character string indicating the datatype of the input data or error message if invalid.
#'   Possible values are "long", "wide", "list" and "upset".
#'   "long" indicates the data is in long format.
#'   "wide" indicates the data is in wide format.
#'   "list" indicates the data is a list.
#'   "upset" indicates the data is a UpsetPlotData object.
detect_upset_datatype <- function(data, group_by = NULL, id_by = NULL) {
    if (inherits(data, "UpsetPlotData")) {
        return("upset")
    }

    if (inherits(data, "data.frame")) {
        if (length(group_by) < 2 && !is.null(group_by)) {
            return("long")
        } else { # length(group_by) > 1
            return("wide")
        }
    }

    if (inherits(data, "list")) {
        return("list")
    }

    stop("Invalid data type. Please provide a data frame, a list or an UpsetPlotData object generated by prepare_upset_data().")
}


#' Prepare data for Upset plot
#'
#' @param data A data frame or a list or an UpsetPlotData object.
#' @param in_form A character string indicating the datatype of the input data.
#'   Possible values are "long", "wide", "list", "upset" or "auto".
#'   "long" indicates the data is in long format.
#'   "wide" indicates the data is in wide format.
#'   "list" indicates the data is a list.
#'   "upset" indicates the data is a UpsetPlotData object.
#'   "auto" indicates the function will detect the datatype of the input data.
#'
#' A long format data would look like:
#' \preformatted{
#' group_by id_by
#' A        a1
#' A        a2
#' B        a1
#' B        a3
#' ...
#' }
#'
#' A wide format data would look like:
#' \preformatted{
#' A    B
#' TRUE TRUE
#' TRUE FALSE
#' FALSE TRUE
#' ...
#' }
#'
#' A list format data would look like:
#' \preformatted{
#' list(A = c("a1", "a2"), B = c("a1", "a3"))
#' }
#'
#' An UpsetPlotData object is generated by prepare_update_data() would look like:
#' \preformatted{
#' group_by
#' --------
#' list("A")  # a2
#' list("B")  # a3
#' list(c("A", "B"))  # a1
#' ...
#' }
#'
#' @param group_by A character string specifying the column name of the data frame to group the data.
#' @param group_by_sep A character string to concatenate the columns in `group_by`,
#'   if multiple columns are provided and the in_form is "long".
#' @param id_by A character string specifying the column name of the data frame to identify the instances.
#'  Required when `group_by` is a single column and data is a data frame.
#' @param specific A logical value to show the specific intersections only.
#'  ggVennDiagram, by default, only return the specific subsets of a region. However,
#'  sometimes, we want to show all the overlapping items for two or more sets.
#'  See \url{https://github.com/gaospecial/ggVennDiagram/issues/64} for more details.
#' @return A UpsetPlotData object
#' @keywords internal
#' @importFrom rlang sym
#' @importFrom dplyr distinct %>%
#' @importFrom utils getFromNamespace
#' @importFrom tidyr uncount
prepare_upset_data <- function(data, in_form = "auto", group_by = NULL, group_by_sep = "_", id_by = NULL, specific = TRUE) {
    if (in_form == "auto") {
        in_form <- detect_upset_datatype(data, group_by, id_by)
    }

    if (in_form == "upset") {
        if (!is.null(group_by)) {
            warning("The group_by argument is ignored when the input data is an UpsetPlotData object.", immediate. = TRUE)
        }
        return(data)
    }

    process_upset_data <- getFromNamespace("process_upset_data", "ggVennDiagram")
    if (in_form == "list") {
        listdata <- data
        group_order <- names(listdata)
    } else if (in_form == "long") {
        group_by <- check_columns(data, group_by, force_factor = TRUE, allow_multi = TRUE, concat_multi = TRUE, concat_sep = group_by_sep)
        group_order <- levels(data[[group_by]])
        listdata <- split(data[[id_by]], data[[group_by]])
    } else { # in_form == "wide"
        group_by <- check_columns(data, group_by, allow_multi = TRUE)
        if (is.null(group_by)) {
            group_by <- setdiff(colnames(data), id_by)
        }
        group_order <- group_by
        for (g in group_by) {
            # columns must be logical or 0/1
            if (!is.logical(data[[g]]) && !all(data[[g]] %in% c(0, 1))) {
                stop("[UpSetPlot] The columns in group_by must be logical or 0/1 when the in_form is 'wide', but column '", g, "' is not.")
            }
        }
        id_by <- check_columns(data, id_by)
        if (is.null(id_by)) {
            id_by <- ".id"
            data$.id <- paste0("id", seq_len(nrow(data)))
        } else {
            # check if id_by column values are unique
            if (any(duplicated(data[[id_by]]))) {
                stop("[UpSetPlot] The values in id_by column must be unique when the in_form is 'wide', but there are duplicated values in column '", id_by, "'.")
            }
        }
        listdata <- lapply(group_by, function(g) data[as.logical(data[[g]]), id_by, drop = TRUE])
        names(listdata) <- group_by
        for (nm in names(listdata)) {
            if (length(listdata[[nm]]) == 0) {
                warning("The group '", nm, "' has no elements, ignored.", immediate. = TRUE)
                listdata[[nm]] <- NULL
            }
        }
    }

    data <- process_upset_data(ggVennDiagram::Venn(listdata), specific = specific)
    idnames <- paste0(as.character(data$left_data$set), " (", as.character(data$left_data$size), ")")
    names(idnames) <- as.character(seq_along(idnames))

    sep <- ifelse(specific, "/", "~")
    data <- distinct(data$main_data, !!sym("id"), !!sym("size"))

    data$Intersection <- lapply(as.character(data$id), function(x) {
        idnames[strsplit(x, sep, fixed = TRUE)[[1]]]
    })

    out <- uncount(data, !!sym("size"))
    attr(out, "group_order") <- group_order
    out
}

#' Atomic Upset plot
#'
#' @inheritParams common_args
#' @inheritParams prepare_upset_data
#' @param label A logical value to show the labels on the bars.
#' @param label_fg A character string specifying the color of the label text.
#' @param label_size A numeric value specifying the size of the label text.
#' @param label_bg A character string specifying the background color of the label.
#' @param label_bg_r A numeric value specifying the radius of the background of the label.
#' @param combmatrix_gap A numeric value specifying the gap between the rows of the combination matrix.
#' The default value is 6, which is suitable for a base_size of 12.
#' The actual gap will be scaled by the text size scale, which is calculated as base_size / 12.
#' @param ... Additional arguments passed to [ggupset::scale_x_upset].
#' @return A ggplot object with Upset plot
#' @keywords internal
#' @importFrom rlang %||% sym
#' @importFrom ggplot2 geom_bar labs guide_colorbar scale_fill_gradientn
#' @importFrom ggplot2 aes geom_rect geom_point geom_line scale_x_continuous expansion
#' @importFrom ggplot2 scale_fill_manual scale_color_manual guides element_blank unit
#' @importFrom ggrepel geom_text_repel
UpsetPlotAtomic <- function(
    data, in_form = "auto", group_by = NULL, group_by_sep = "_", id_by = NULL,
    label = TRUE, label_fg = "black", label_size = NULL, label_bg = "white", label_bg_r = 0.1,
    palette = "material-indigo", palcolor = NULL, alpha = 1, specific = TRUE, combmatrix_gap = 6,
    theme = "theme_this", theme_args = list(), title = NULL, subtitle = NULL, xlab = NULL, ylab = NULL,
    aspect.ratio = 0.6, legend.position = "right", legend.direction = "vertical", levels = NULL, ...) {
    ggplot <- if (getOption("plotthis.gglogger.enabled", FALSE)) {
        gglogger::ggplot
    } else {
        ggplot2::ggplot
    }

    data <- prepare_upset_data(data, in_form, group_by, group_by_sep, id_by, specific)
    group_order <- attr(data, "group_order")

    base_size <- theme_args$base_size %||% 12
    text_size_scale <- base_size / 12

    p <- ggplot(data, aes(x = !!sym("Intersection"))) +
        geom_bar(aes(fill = after_stat(!!sym("count"))), alpha = alpha, color = "black", width = 0.5) +
        scale_fill_gradientn(
            n.breaks = 3,
            colors = palette_this(palette = palette, palcolor = palcolor),
            na.value = "grey80",
            guide = guide_colorbar(
                title = "", alpha = alpha,
                frame.colour = "black", ticks.colour = "black", title.hjust = 0
            )
        )
    if (isTRUE(label)) {
        p <- p + geom_text_repel(aes(label = after_stat(!!sym("count"))),
            stat = "count",
            colour = label_fg, size = label_size %||% text_size_scale * 3.5,
            bg.color = label_bg, bg.r = label_bg_r, nudge_y = 0.08 * text_size_scale,
            point.size = NA, max.overlaps = 100, force = 0,
            min.segment.length = 0, segment.colour = NA
        )
    }
    upset_args <- list(...)
    ytrans <- upset_args$ytrans %||% "identity"

    # scale_x_upset() internally calls axis_combmatrix(), bundling a CoordCombMatrix into its
    # return list. We extract only the ScaleUpset (index [[1]]) and discard the bundled
    # CoordCombMatrix (index [[2]]), then add our own axis_combmatrix() once — avoiding
    # a second coord system replacement and its "Coordinate system already present" message.
    scale_x_upset_res <- ggupset::scale_x_upset(...)
    p <- p +
        labs(title = title, subtitle = subtitle, x = xlab %||% "", y = ylab %||% "Intersection size") +
        scale_x_upset_res[[1]] +
        do.call(theme, theme_args) +
        ggplot2::theme(
            aspect.ratio = aspect.ratio,
            legend.position = legend.position,
            legend.direction = legend.direction,
            panel.grid.major = element_line(colour = "grey80", linetype = 2)
        ) +
        # try to keep the group order in combmatrix
        ggupset::axis_combmatrix(sep = " // ", ytrans = ytrans, override_plotting_function = function(df) {
            # df$single_label <- factor(df$single_label, levels = group_order)
            # single_label is something like: 'C (5)''B (5)''D (5)''A (5)'
            # but group_order is something like: 'A', 'B', 'C', 'D'
            # so we need to extract the group name from single_label and match it with group_order
            df$group_name <- sapply(as.character(df$single_label), function(x) {
                # extract the group name from single_label, which is the part before the first " ("
                strsplit(x, " (", fixed = TRUE)[[1]][1]
            })
            df$group_name <- factor(df$group_name, levels = rev(group_order))
            df <- df[order(df$group_name), , drop = FALSE]
            df$single_label <- factor(df$single_label, levels = unique(df$single_label))
            ggplot(df, aes(x = !!sym("at"), y = !!sym("single_label"))) +
                geom_rect(
                    aes(fill = !!sym("index") %% 2 == 0), ymin = df$index - 0.5,
                        ymax = df$index + 0.5, xmin = 0, xmax = 1) +
                geom_point(aes(color = !!sym("observed")), size = text_size_scale * 3) +
                geom_line(data = function(dat) dat[dat$observed, , drop = FALSE],
                    aes(group = !!sym("labels")), linewidth = 1.2 * text_size_scale) +
                ggplot2::ylab("") + ggplot2::xlab("") +
                scale_x_continuous(limits = c(0, 1), expand = c(0, 0)) +
                scale_y_discrete(expand = c(0, 0)) +
                scale_fill_manual(values= c(`TRUE` = "white", `FALSE` = "#F7F7F7")) +
                scale_color_manual(values= c(`TRUE` = "black", `FALSE` = "#E0E0E0")) +
                guides(color = "none", fill = "none") +
                ggplot2::theme(
                    panel.background = element_blank(),
                    axis.text.x = element_blank(),
                    axis.text.y = element_text(size = 12 * text_size_scale^0.7, color = "black"),
                    axis.ticks.y = element_blank(),
                    axis.ticks.length = unit(0, "pt"),
                    axis.title.y = element_blank(),
                    axis.title.x = element_blank(),
                    axis.line = element_blank(),
                    panel.border = element_blank()
                )
        }) +
        ggupset::theme_combmatrix(combmatrix.label.extra_spacing = combmatrix_gap * text_size_scale)

    n_sets <- upset_args$n_sets %||% 99
    n_sets <- min(n_sets, length(unique(unlist(data$Intersection))))
    n_intersections <- upset_args$n_intersections %||% 99
    n_intersections <- min(n_intersections, length(unique(data$Intersection)))
    maxchars <- max(sapply(unique(unlist(data$Intersection)), nchar))

    # Height driven by number of sets; width driven by number of intersections.
    # x_scale_factor = 0.6 preserves the original per-intersection width (the original
    # default aspect.ratio = 0.6 was used as a width-per-intersection scale, not H/W ratio).
    # aspect.ratio is now used as the standard H/W ratio coupling.
    dims <- calculate_plot_dimensions(
        base_height = 4.5 + n_sets * 0.5,
        aspect.ratio = aspect.ratio,
        n_x = n_intersections,
        x_scale_factor = 0.6,
        legend.position = legend.position,
        legend.direction = legend.direction,
        legend_n = n_sets,
        legend_nchar = maxchars
    )

    attr(p, "height") <- dims$height
    attr(p, "width") <- dims$width + maxchars * 0.05

    p
}

#' Upset Plot
#'
#' @inheritParams common_args
#' @inheritParams UpsetPlotAtomic
#' @return A ggplot object or wrap_plots object or a list of ggplot objects
#' @export
#' @examples
#' \donttest{
#' data <- list(
#'     A = 1:5,
#'     B = 2:6,
#'     C = 3:7,
#'     D = 4:8
#' )
#' UpsetPlot(data)
#' UpsetPlot(data, label = FALSE)
#' UpsetPlot(data, palette = "Reds", specific = FALSE)
#'
#' # long form input
#' data_long <- data.frame(
#'     group_by = factor(
#'          c(rep("A", 5), rep("B", 5), rep("C", 5), rep("D", 5)),
#'          levels = c("A", "B", "C", "D")
#'     ),
#'     id_by = c(1:5, 2:6, 3:7, 4:8)
#' )
#' UpsetPlot(data_long, in_form = "long", group_by = "group_by", id_by = "id_by")
#'
#' # wide form input
#' data <- data.frame(
#'     id = LETTERS[1:10],
#'     B = c(1, 0, 1, 1, 0, 0, 1, 0, 1, 0),
#'     A = c(1, 1, 1, 0, 0, 1, 0, 0, 1, 0),
#'     D = c(1, 0, 0, 1, 1, 0, 0, 1, 0, 1),
#'     C = c(0, 1, 1, 0, 1, 0, 1, 0, 1, 0)
#' )
#' UpsetPlot(data, in_form = "wide", id_by = "id", n_intersections = 4)
#' }
UpsetPlot <- function(
    data, in_form = c("auto", "long", "wide", "list", "upset"), split_by = NULL, split_by_sep = "_",
    group_by = NULL, group_by_sep = "_", id_by = NULL, label = TRUE, label_fg = "black",
    label_size = NULL, label_bg = "white", label_bg_r = 0.1, palette = "Blues", palcolor = NULL,
    alpha = 1, specific = TRUE, theme = "theme_this", theme_args = list(), title = NULL, subtitle = NULL,
    xlab = NULL, ylab = NULL, aspect.ratio = 0.6, legend.position = "right", legend.direction = "vertical",
    combine = TRUE, nrow = NULL, ncol = NULL, byrow = TRUE, seed = 8525, combmatrix_gap = 6,
    axes = NULL, axis_titles = axes, guides = NULL, design = NULL, ...) {
    validate_common_args(seed)
    in_form <- match.arg(in_form)
    theme <- process_theme(theme)
    if (!is.null(split_by) && !inherits(data, "data.frame")) {
        stop("'split_by' is only available for data frame.")
    }
    split_by <- check_columns(data, split_by, force_factor = TRUE, allow_multi = TRUE, concat_multi = TRUE, concat_sep = split_by_sep)

    if (!is.null(split_by)) {
        data[[split_by]] <- droplevels(data[[split_by]])
        datas <- split(data, data[[split_by]])
        # keep the order of levels
        datas <- datas[levels(data[[split_by]])]
    } else {
        datas <- list(data)
        names(datas) <- "..."
    }
    palette <- check_palette(palette, names(datas))
    palcolor <- check_palcolor(palcolor, names(datas))
    legend.direction <- check_legend(legend.direction, names(datas), "legend.direction")
    legend.position <- check_legend(legend.position, names(datas), "legend.position")

    plots <- lapply(
        names(datas), function(nm) {
            default_title <- if (length(datas) == 1 && identical(nm, "...")) NULL else nm
            if (is.function(title)) {
                title <- title(default_title)
            } else {
                title <- title %||% default_title
            }
            if (is.null(group_by) && (in_form %in% c("auto", "wide"))) {
                group_by <- setdiff(colnames(datas[[nm]]), c(id_by, split_by))
            }
            UpsetPlotAtomic(datas[[nm]],
                in_form = in_form, group_by = group_by, group_by_sep = group_by_sep, id_by = id_by, combmatrix_gap = combmatrix_gap,
                label = label, label_fg = label_fg, label_size = label_size, label_bg = label_bg, label_bg_r = label_bg_r,
                palette = palette[[nm]], palcolor = palcolor[[nm]], alpha = alpha, specific = specific,
                theme = theme, theme_args = theme_args, title = title, subtitle = subtitle, xlab = xlab, ylab = ylab,
                aspect.ratio = aspect.ratio, legend.position = legend.position[[nm]], legend.direction = legend.direction[[nm]],
                ...
            )
        }
    )

    combine_plots(plots, combine = combine, nrow = nrow, ncol = ncol, byrow = byrow,
        axes = axes, axis_titles = axis_titles, guides = guides, design = design)
}
