#' @importFrom R6 R6Class
NULL

#' An R6 Class for Generating Theseus Plot
#'
#' @description
#' The \code{ShipOfTheseus} class decomposes the difference in outcome rates
#' between two datasets. For a selected column, it computes subgroup
#' contributions, summarizes the results in tables, and visualizes them as
#' waterfall-style Theseus Plots.
#'
#' @import dplyr ggplot2 stringr
#'
#' @export
ShipOfTheseus <- R6::R6Class(
  "ShipOfTheseus",

  private = list(
    labels = NULL,
    compute_scores = NULL,
    to_factor = NULL,
    compute_contribution = NULL,
    compute_info = NULL,
    compute_size = NULL,
    xlab = NULL,
    ylab = NULL,
    digits = NULL,
    text_size = NULL,
    
    #' @importFrom forcats fct_na_value_to_level
    #' @importFrom tidyr replace_na
    prepare_input_data = function(data, outcome) {
      data |>
        # Fill missing values while preserving column types, since downstream
        # table()/plot() handle character and factor columns differently
        mutate(
          across(where(is.character), \(x) replace_na(x, "(Missing)")),
          across(where(is.factor) & where(anyNA), \(x) fct_na_value_to_level(x, level = "(Missing)"))
        ) |>
        # Standardize the outcome column name for internal use
        rename(.outcome = all_of(outcome))
    },
    
    prepare_plot_data = function(column_name, n, levels, continuous, flip) {
      df_contrib <- self$table(!!rlang::sym(column_name), n = n, continuous = continuous) |>
        select(items = 1L, contrib)
      
      if (flip) {
        df_contrib <- df_contrib |> mutate(contrib = -contrib)
      }
      
      is_factor <- is.factor(df_contrib$items)
      
      if (is_factor) {
        df_contrib <- if (flip) arrange(df_contrib, desc(items)) else arrange(df_contrib, items)
      } else {
        df_contrib <- arrange(df_contrib, contrib)
      }
      
      data_size <- private$compute_size(
        column_name,
        target = df_contrib$items,
        continuous = continuous
      )
      
      if (!is.null(levels)) {
        levels <- as.character(levels)
        if (flip) levels <- rev(levels)
        df_contrib <- data.frame(items = levels) |> inner_join(df_contrib, by = "items")
      }
      
      list(
        df_contrib = df_contrib,
        data_size = data_size
      )
    },
    
    compute_bar_height = function(data_waterfall, data_size, max_score, main_item = NULL,
                                  bar_max_value = NULL, bar_min_ratio = 0.15) {
      if (is.null(main_item) & is.null(bar_max_value)) {
        max_amount <- max_score * 100 * bar_min_ratio
        n_max <- data_size |> tail(-1) |> 
          slice_max(order_by = n, n = 1, with_ties = FALSE) |> pull(n)
      } else if(!is.null(main_item)) {
        max_amount <- data_waterfall |> filter(items == main_item) |> pull(contrib) |> abs()
        n_max <- data_size |> filter(items == main_item) |> pull(n) |> max()
      } else if(!is.null(bar_max_value)) {
        max_amount <- bar_max_value
        n_max <- data_size |> filter(n == max(n)) |> pull(n) |> max()
      }
      max_amount / n_max
    }
  ),

  public = list(
    #' @description
    #' The constructor of the \code{ShipOfTheseus} class.
    #'
    #' @param data1 data frame representing the first group (e.g., the baseline
    #'   data).
    #' @param data2 data frame representing the second group (e.g., the
    #'   comparison data).
    #' @param outcome string specifying the outcome variable used to compute the
    #'   rate metric (default is \code{"y"}). Typically, this is a binary
    #'   indicator (e.g., 0/1) that is aggregated to form rates.
    #' @param labels character vector of length 2 giving the labels for the two
    #'   groups. The first corresponds to \code{data1}, the second to
    #'   \code{data2}. Default is \code{c("Baseline", "Comparison")}.
    #' @param xlab string specifying the x-axis label for plots. If \code{NULL}
    #'   (default), no label is displayed.
    #' @param ylab string specifying the y-axis label for plots. If \code{NULL}
    #'   (default), no label is displayed.
    #' @param digits integer indicating the number of decimal places to use for
    #'   displaying numeric values (default is \code{1}).
    #' @param text_size numeric value specifying the relative size of text
    #'   elements in plots (default is \code{1.0}).
    #'
    #' @return A \code{ShipOfTheseus} object, which can be used with
    #'   \code{plot()} to create Theseus plots.
    #'
    #' @importFrom forcats fct_c fct_na_value_to_level
    #' @importFrom memoise memoise
    #' @importFrom tibble tibble
    #' @importFrom tidyr replace_na
    initialize = function(data1, data2, outcome, labels, xlab, ylab, digits, text_size) {
      outcome <- rlang::quo_squash(outcome) |> rlang::as_string()
      labels <- as.character(labels)

      data1 <- private$prepare_input_data(data1, outcome)
      data2 <- private$prepare_input_data(data2, outcome)

      private$labels <- labels
      private$xlab <- xlab 
      private$ylab <- ylab
      private$digits <- digits
      private$text_size <- text_size

      private$compute_scores <- memoise(function(column_name) {
        score1 <- data1 |> summarise(score = mean(.outcome)) |> pull(score)
        score2 <- data2 |> summarise(score = mean(.outcome)) |> pull(score)
        c(score1, score2)
      })

      private$to_factor <- memoise(function(column_name, continuous) {
        if (is.null(continuous$breaks)) {
          values <- c(data1[[column_name]], data2[[column_name]])
          break_num <- continuous$n
          if (continuous$split == "width") {
            if (any(is.na(values))) break_num <- break_num - 1L
            min <- min(values, na.rm = TRUE)
            max <- max(values, na.rm = TRUE)
            breaks <- seq(min, max, length.out = break_num + 1)
          } else if(continuous$split == "count") {
            breaks <- compute_breaks(values, break_num = break_num)
          } else {
            breaks <- compute_breaks(values, break_num = break_num * 20L)
            breaks <- sort(unique(breaks))
            d1 <- data1 |>
              select(x = !!rlang::sym(column_name), y = .outcome) |>
              filter(!is.na(x))
            d2 <- data2 |>
              select(x = !!rlang::sym(column_name), y = .outcome) |>
              filter(!is.na(x))
            while (length(breaks) > break_num + 1L) {
              data1_tmp <- d1 |>
                mutate(x = cut(x, breaks = breaks, include.lowest = TRUE)) |>
                group_by(x) |>
                summarise(y = mean(y)) |>
                mutate(diff1 = abs(lead(y) - y))
              data2_tmp <- d2 |>
                mutate(x = cut(x, breaks = breaks, include.lowest = TRUE)) |>
                group_by(x) |>
                summarise(y = mean(y)) |>
                mutate(diff2 = abs(lead(y) - y))
              data_tmp <- data1_tmp |> left_join(data2_tmp, by = "x") |>
                mutate(diff = sqrt(diff1^2 + diff2^2))
              breaks <- breaks[-(which.min(data_tmp$diff) + 1L)]
            }
          }
          if (continuous$pretty) {
            breaks <- pretty_breaks(breaks)
            if (any(table(breaks) >= 2)) {
              warning("Prettying breaks reduced the number of breaks. Try pretty = FALSE.")
              breaks <- unique(breaks)
            }
          }
        } else {
          breaks <- continuous$breaks
        }

        df1 <- data1
        df2 <- data2

        df1[[column_name]] <- cut(df1[[column_name]], breaks = breaks, include.lowest = TRUE, dig.lab = 50)
        df2[[column_name]] <- cut(df2[[column_name]], breaks = breaks, include.lowest = TRUE, dig.lab = 50)

        df1[[column_name]] <- fct_na_value_to_level(df1[[column_name]], level = "(Missing)")
        df2[[column_name]] <- fct_na_value_to_level(df2[[column_name]], level = "(Missing)")

        list(df1, df2)
      })

      private$compute_contribution <- memoise(function(column_name, continuous) {
        if (is.numeric(data1[[column_name]])) {
          data_list <- private$to_factor(column_name, continuous)
          data1 <- data_list[[1]]
          data2 <- data_list[[2]]
        }

        df1 <- data1 |>
          group_by(!!rlang::sym(column_name)) |>
          summarise(y = sum(.outcome), n = n(), rate = y / n)
        df2 <- data2 |>
          group_by(!!rlang::sym(column_name)) |>
          summarise(y = sum(.outcome), n = n(), rate = y / n)

        names1 <- df1[[column_name]]
        names2 <- df2[[column_name]]

        scores <- private$compute_scores(column_name)
        score1 <- scores[1]
        score2 <- scores[2]

        df_contrib <- tibble()
        for (name in names2) {
          df_temp <- df1
          if (name %in% names1) {
            df_temp[df_temp[[column_name]] == name, ] <- df2[df2[[column_name]] == name, ]
          } else {
            df_temp <- rbind(df_temp, df2[df2[[column_name]] == name, ])
          }

          score_new <- df_temp |> summarise(score = sum(y) / sum(n)) |> pull(score)
          diff <- score_new - score1
          res <- tibble(items = name, amount = diff)
          df_contrib <- rbind(df_contrib, res)
        }
        for (name in names1) {
          df_temp <- df2
          if (name %in% names2) {
            df_temp[df_temp[[column_name]] == name, ] <- df1[df1[[column_name]] == name, ]
          } else {
            df_temp <- rbind(df_temp, df1[df1[[column_name]] == name, ])
          }

          score_new <- df_temp |> summarise(score = sum(y) / sum(n)) |> pull(score)
          diff <- score2 - score_new
          res <- tibble(items = name, amount = diff)
          df_contrib <- rbind(df_contrib, res)
        }

        if (is.factor(names1)) {
          names <- fct_c(names1, names2)
          df_contrib <- df_contrib |>
            mutate(items = factor(items, levels = levels(names)))
        }
        df_contrib |>
          group_by(items) |>
          summarise(contrib = mean(amount)) |>
          mutate(contrib = (score2 - score1) * contrib / sum(contrib))
      })

      private$compute_info <- memoise(function(column_name, continuous) {
        if (is.numeric(data1[[column_name]])) {
          data_list <- private$to_factor(column_name, continuous)
          data1 <- data_list[[1]]
          data2 <- data_list[[2]]
        }

        data1_info <- data1 |>
          group_by(items = !!rlang::sym(column_name)) |>
          summarise(n1 = n(), x1 = sum(.outcome), rate1 = x1 / n1)
        data2_info <- data2 |>
          group_by(items = !!rlang::sym(column_name)) |>
          summarise(n2 = n(), x2 = sum(.outcome), rate2 = x2 / n2)
        data1_info |> full_join(data2_info, by = "items") |>
          select(items, starts_with("n"), starts_with("x"), starts_with("rate")) |>
          replace_na(list(n1 = 0L, n2 = 0L, x1 = 0L, x2 = 0L))
      })

      private$compute_size <- memoise(function(column_name, target, continuous) {
        if (is.numeric(data1[[column_name]])) {
          data_list <- private$to_factor(column_name, continuous)
          data1 <- data_list[[1]]
          data2 <- data_list[[2]]
        }

        data1_size <- data1 |>
          filter(!!rlang::sym(column_name) %in% target) |>
          count(items = !!rlang::sym(column_name)) |>
          mutate(type = labels[1])
        data2_size <- data2 |>
          filter(!!rlang::sym(column_name) %in% target) |>
          count(items = !!rlang::sym(column_name)) |>
          mutate(type = labels[2])
        item_names <- unique(c(data1_size$items, data2_size$items))
        other_name <- target[!(target %in% item_names)]

        if (length(other_name) == 0L) {
          bind_rows(data1_size, data2_size)
        } else {
          data1_size_other <- data1 |>
            filter(!(!!rlang::sym(column_name) %in% target)) |>
            count() |>
            mutate(type = labels[1], items = other_name)
          data2_size_other <- data2 |>
            filter(!(!!rlang::sym(column_name) %in% target)) |>
            count() |>
            mutate(type = labels[2], items = other_name)
          if (is.factor(data1_size$items)) {
            data1_size <- data1_size |> mutate(items = as.character(items))
            data2_size <- data2_size |> mutate(items = as.character(items))
            bind_rows(data1_size, data1_size_other, data2_size, data2_size_other)
          } else {
            bind_rows(data1_size, data1_size_other, data2_size, data2_size_other)
          }
        }
      })
    },

    #' @description
    #' Generate a contribution table for a given column.
    #'
    #' @param column_name string. The name of the column to analyze.
    #' @param n integer. Maximum number of top contributing subgroups to display.
    #'   If the number of subgroups exceeds `n`, the remaining are aggregated.
    #' @param continuous list. A configuration list for handling continuous
    #'   variables (e.g., specifying number of bins or custom breaks).
    #'
    #' @return A tibble summarizing subgroup contributions to the difference
    #'   between the two groups, including counts, total outcomes, and rates
    #'   for each subgroup.
    table = function(column_name, n = Inf, continuous = continuous_config()) {
      column_name <- rlang::ensym(column_name) |> rlang::as_string()
      data_contrib <- private$compute_contribution(column_name, continuous)
      data_info <- private$compute_info(column_name, continuous)
      result <- data_contrib |>
        left_join(data_info, by = "items")
      is_factor <- is.factor(result$items)
      if (is_factor) {
        levels <- levels(result$items)
        result <- result |> arrange(items)
      } else {
        result <- result |> arrange(desc(abs(contrib)))
      }
      n_items <- nrow(result)
      if (n_items > n) {
        n_other <- n_items - n + 1L
        result_head <- head(result |> arrange(desc(abs(contrib))), n - 1L) |>
          mutate(items = as.character(items))
        result_tail <- tail(result |> arrange(desc(abs(contrib))), n_other)
        result_other <- result_tail |>
          mutate(items = str_glue("Sum of {n_other} other subgroups")) |>
          group_by(items) |>
          summarise_at(vars(contrib, n1, n2, x1, x2), sum) |>
          mutate(rate1 = x1 / n1, rate2 = x2 / n2)
        result <- bind_rows(result_head, result_other)
        if (is_factor) {
          levels <- c(levels, str_glue("Sum of {n_other} other subgroups"))
          result <- result |>
            mutate(items = factor(items, levels = levels)) |>
            arrange(items)
        }
      }
      names(result)[1] <- column_name
      result
    },

    #' @description
    #' Generate a Theseus plot for a specified column
    #'
    #' @param column_name The name of the column to visualize.
    #' @param n integer. Maximum number of top contributing subgroups to display.
    #'   Remaining subgroups are aggregated if necessary.
    #' @param main_item string. The subgroup used as the reference for scaling
    #'   the bar heights.
    #' @param bar_max_value numeric. Maximum value for scaling the contribution bars.
    #' @param levels character vector specifying the display order of subgroups.
    #' @param continuous list. Configuration for handling continuous variables
    #'   (e.g., number of bins or custom breaks).
    #'
    #' @return A ggplot object representing the Theseus Plot for the specified column.
    #'
    #' @importFrom tibble tibble
    #' @importFrom waterfalls waterfall
    plot = function(column_name, n = 10L, main_item = NULL, bar_max_value = NULL,
                    levels = NULL, continuous = continuous_config()) {
      column_name <- rlang::ensym(column_name) |> rlang::as_string()

      labels <- private$labels
      scores <- private$compute_scores(column_name)
      score1 <- scores[1]
      max_score <- max(scores)
      plot_data <- private$prepare_plot_data(column_name, n, levels, continuous, flip = FALSE)
      
      df_contrib <- plot_data$df_contrib
      data_size <- plot_data$data_size

      names <- as.character(df_contrib$items)
      data_waterfall <- tibble(items = labels[1], contrib = score1) |>
        bind_rows(df_contrib)|>
        mutate(contrib = round(contrib * 100, digits = private$digits))

      n_bars <- nrow(data_waterfall) + 1L
      p <- waterfall(
        data_waterfall,
        calc_total = TRUE, total_axis_text = n_bars,
        total_rect_text_color = "black", total_rect_color = "#00BFC4", 
        rect_text_size = private$text_size)
      scale_x <- p@scales$get_scales("x")
      scale_x$labels <- c(data_waterfall$items, labels[2])
      
      bar_rate <- private$compute_bar_height(data_waterfall, data_size, max_score, 
                                             main_item, bar_max_value)

      data_size <- tibble(x = factor(seq_len(n_bars))) |>
        mutate(items = c(labels[1], names, labels[2])) |>
        left_join(data_size, by = "items") |>
        tidyr::replace_na(list(n = 0L)) |>
        mutate(type = factor(type, levels = labels)) |>
        mutate(n = n * bar_rate)
            
      p <- p +
        geom_col(data = data_size, aes(x, n, fill = type), width = 0.7, position = position_dodge()) +
        scale_fill_manual(values = c("#7CAE00", "#C77CFF"), guide = "none")
      p$layers <- append(head(p$layers, -1), tail(p$layers, 1), 1)
      
      current_theme <- ggplot2::theme_get()
      current_theme$text <- element_text(size = 11 * private$text_size)
      p + current_theme + xlab(private$xlab) + ylab(private$ylab)
    },

    #' @description
    #' Generate a Theseus plot for a specified column
    #'
    #' @param column_name The name of the column to visualize.
    #' @param n integer. Maximum number of top contributing subgroups to display.
    #'   Remaining subgroups are aggregated if necessary.
    #' @param main_item string. The subgroup used as the reference for scaling
    #'   the bar heights.
    #' @param bar_max_value numeric. Maximum value for scaling the contribution bars.
    #' @param levels character vector specifying the display order of subgroups
    #' @param continuous list. Configuration for handling continuous variables
    #'   (e.g., number of bins or custom breaks).
    #'
    #' @return A ggplot object representing the Theseus Plot for the specified column.
    #'
    #' @importFrom tidyr replace_na
    #' @importFrom tibble tibble
    #' @importFrom waterfalls waterfall
    plot_flip = function(column_name, n = 10L, main_item = NULL, bar_max_value = NULL,
                         levels = NULL, continuous = continuous_config()) {
      column_name <- rlang::ensym(column_name) |> rlang::as_string()

      labels <- private$labels
      scores <- private$compute_scores(column_name)
      score2 <- scores[2]
      max_score <- max(scores)
      plot_data <- private$prepare_plot_data(column_name, n, levels, continuous, flip = TRUE)
      
      df_contrib <- plot_data$df_contrib
      data_size <- plot_data$data_size
      
      names <- as.character(df_contrib$items)
      data_waterfall <- tibble::tibble(items = labels[2], contrib = score2) |>
        bind_rows(df_contrib)|>
        mutate(contrib = round(contrib * 100, digits = private$digits))

      colors <- if_else(data_waterfall$contrib > 0, "#F8766D", "#00BFC4")
      colors[1] <- "#00BFC4"
      n_bars <- nrow(data_waterfall) + 1L
      p <- waterfall(
        data_waterfall,
        calc_total = TRUE, total_axis_text = n_bars,
        total_rect_text_color = "black", fill_colours = colors,
        fill_by_sign = FALSE, total_rect_color = "#00BFC4",
        rect_text_size = private$text_size) +
        coord_flip()
      scale_x <- p@scales$get_scales("x")
      scale_x$labels <- c(data_waterfall$items, labels[1])
      
      reverse_sign <- function(x) {
        x <- str_replace(x, "\u2212", "-")
        x <- -as.numeric(x)
        x <- as.character(x)
        str_replace(x, "-", "\u2212")
      }

      is_after <- TRUE
      for (i in seq_along(p$layers |> head(-2))) {
        if ("GeomText" %in% class(p$layers[[i]]$geom)) {
          if (is_after) {
            is_after <- FALSE
            next
          }
          p$layers[[i]]$aes_params$label <- reverse_sign(p$layers[[i]]$aes_params$label)
        }
      }

      bar_rate <- private$compute_bar_height(data_waterfall, data_size, max_score, 
                                             main_item, bar_max_value)
      
      data_size <- tibble(x = factor(seq_len(n_bars))) |>
        mutate(items = c(labels[2], names, labels[1])) |>
        left_join(data_size, by = "items") |>
        tidyr::replace_na(list(n = 0L)) |>
        mutate(type = factor(type, levels = rev(labels))) |>
        mutate(n = n * bar_rate)
      
      p <- p +
        geom_col(data = data_size, aes(x, n, fill = type), width = 0.7, position = position_dodge()) +
        scale_fill_manual(values = c("#C77CFF", "#7CAE00"), guide = "none")
      p$layers <- append(head(p$layers, -1), tail(p$layers, 1), 1)
      
      current_theme <- ggplot2::theme_get()
      current_theme$text <- element_text(size = 11 * private$text_size)
      p + current_theme + xlab(private$xlab) + ylab(private$ylab)
    }
  )
)
