#' Theme for ggplot2
#' @param base_size Legend title size, all other sizes scaled appropriately 
#' to this
#' @param base_family Font family
#' 
#' @import ggplot2
#' 
theme_bma <- function(
  base_size = 10, 
  base_family = get_base_family()
) {

  `%+replace%` <- ggplot2::`%+replace%`

  base_family <- ifelse(is.null(base_family), get_base_family(), base_family)

  ggplot2::theme_bw(base_size = base_size, base_family = base_family) %+replace%
    ggplot2::theme(
      text = ggplot2::element_text(family = base_family),
      plot.title = ggplot2::element_text(
        size = base_size + 6,
        margin = margin(0, 0, 13, 0),
        hjust = 0,
        face = "bold"
      ),
      plot.title.position = "plot",
      plot.subtitle = ggplot2::element_text(
        size = base_size + 4,
        margin = margin(0, 0, 13, 0),
        hjust = 0
      ),
      axis.text = ggplot2::element_text(size = base_size - 1),
      axis.title = ggplot2::element_text(size = base_size + 4),
      axis.title.x = ggplot2::element_text(
        margin = ggplot2::margin(t = base_size - 1)
      ),
      axis.title.y = ggplot2::element_text(
        margin = ggplot2::margin(l = 0, r = base_size + 1), 
        angle = 90
      ),
      legend.text = ggplot2::element_text(size = base_size),
      legend.title = ggplot2::element_text(size = base_size + 2),
      strip.background = ggplot2::element_rect(fill = "grey90")
    )
}


#' Set base font family for ggplot2.
#' 
#' @return Character string of the name of a postscript font related to 
#' Arial if available, otherwise "sans".
#' 
#' @importFrom grDevices postscriptFonts
#' 
get_base_family <- function() {
  avail_fonts <- tryCatch(
    grDevices::postscriptFonts(),
    error = function(cond) NULL,
    warning = function(cond) NULL
  )

  avail_fontnames <- names(avail_fonts)

  if (any(grepl("Arial", avail_fontnames))) {
    base_family <- avail_fontnames[grep("Arial", avail_fontnames)[1]]
  } else {
    base_family <- "sans"
  }

  return(base_family)
}


#' Generate arm closure summaries for batches
#' 
#' @param file_prefix Consistent beginning of filenames holding 
#' arm closure data. Defaults to `closures`.
#' @param run_time Specify a particular instance of `biomkrAccrual()`
#' execution using a date-time format `yyyy-mm-dd-hh-mm-ss`. 
#' Used to select which files will be summarised.
#' @param output_path Directory where the input files are located
#' and the output files will be written.
#' @param keep_files Save data files and plots generated during the run. 
#' Defaults to TRUE.
#' 
#' @export
#' 
#' @importFrom stats sd
#' @importFrom utils read.csv write.csv
#' 
##### Nothing is calling this?
get_arm_closures <- function(
  file_prefix = "closures",
  run_time = "2024-08-07-18-35-09",
  output_path = "../biomkrAccrual_output_data/",
  keep_files = TRUE
) {
  # What output files do we have?
  filenames <- list.files(
    output_path, 
    pattern = paste0("^", file_prefix, "-", run_time, ".*.csv"), 
    full.names = TRUE
  )

  # Read files
  closures_ls <- lapply(filenames, read.csv)

  # Summarise files
  summ <- lapply(closures_ls, function(d) summary(d[, -1]))
  name_types <- c("1", "2", "mixed", "unbalanced", "multirate", "multimix")
  names(summ) <- c(
    paste0("gamma_rate_closures_", name_types),
    paste0("fixed_rate_closures_", name_types)
  )

  write.csv(
    summ, 
    paste0(output_path, "arm_closure_summary-", run_time, ".csv")
  )

  # Standard Deviations
  sd_mx <- t(sapply(
    closures_ls, 
    function(a_df) {
      sapply(
        seq_len(ncol(a_df)),
        function(i) stats::sd(a_df[i, ], na.rm = TRUE)
      )
    }
  ))

  sd_mx <- sd_mx[, -1]

  colnames(sd_mx) <- paste0("T", 1:3, "_sd")
  rownames(sd_mx) <- c(
    paste0("gamma_rate_closures_", name_types),
    paste0("fixed_rate_closures_", name_types)
  )

  write.csv(
    as.data.frame(sd_mx), 
    paste0(output_path, "arm_closures_sd", run_time, ".csv")
  )

  return(list(summ, sd_mx))
}


#' Plot predicted recruitment from file containing a CSV from
#' a single run
#' 
#' @param file_prefix Consistent beginning of filenames holding 
#' arm closure data. Defaults to `accrual`.
#' @param plot_prefix Prefix for file name to identify plot type. 
#' Defaults to `accrual_plot`.
#' @param run_time Specify a particular instance of `biomkrAccrual()`
#' execution using a date-time format `yyyy-mm-dd-hh-mm-ss`. 
#' Used to select which files will be summarised.
#' @param output_path Directory where the output files from the 
#' `biomkrAccrual()` instance are located.
#' @param figs_path Folder where figures generated during execution
#' will be stored; defaults to the `figures` subdirectory in
#' `output_path`.
#' 
#' @export
#' 
accrual_plot_from_file <- function(
  file_prefix = "accrual",
  plot_prefix = "accrual-from-file",
  run_time = "2024-08-07-18-35-09",
  output_path = "../biomkrAccrual_output_data/",
  figs_path = paste0(output_path, "figures/")
) {
  # Validate input

  checkmate::assert_directory_exists(
    file.path(output_path), 
    access = "rx"
  )

  input_file <- paste0(
    output_path, file_prefix, "-", run_time, ".csv"
  )

  checkmate::assert_file_exists(file.path(input_file))

  makeifnot_dir(figs_path, min_access = "rwx")

  accrual_raw_df <- utils::read.csv(file.path(input_file))

  # Get unique arm identifiers
  arm_names <- unique(sapply(
    strsplit(names(accrual_raw_df), "\\."), 
    getElement,
    1
  ))

  # Make dataframe of arm recruitment per week by summing site-arm columns
  accrual_df <- as.data.frame(
    lapply(
      arm_names, 
      function(n) rowSums(accrual_raw_df[startsWith(names(accrual_raw_df), n)])
    ),
    col.names = arm_names
  )

  # Convert to long format of class "accrualplotdata"
  accrual_df <- accrual_to_long(accrual_df)

  # Plot 
  p <- plot(
    accrual_df, 
    plot_prefix = plot_prefix,
    run_time = run_time
  )

  print(p)

  ggplot2::ggsave(
    paste0(figs_path, plot_prefix, "-", run_time, ".png"),
    plot = p,
    width = 12,
    height = 8,
    dpi = 400
  )
}


#' Convert accrual data to long format dataframe of class
#' `accrualplotformat`.
#' 
#' @param accrual_df Wide format accrual data.
#' 
#' @importFrom stats reshape
#' 
accrual_to_long <- function(accrual_df) {
  # Convert to cumulative sums
  accrual_df <- cumsum(accrual_df)

  arm_names <- names(accrual_df)

  # Add week information
  accrual_df$Week <- seq_len(nrow(accrual_df))

  # Convert to long format
  accrual_df <- stats::reshape(
    accrual_df,
    direction = "long",
    varying = list(arm_names),
    timevar = "Arm",
    times = arm_names,
    v.names = "Recruitment",
    idvar = "Week"
  )

  accrual_df$Arm <- factor(accrual_df$Arm)

  # Define an S3 class so we can have a custom plot command
  class(accrual_df) <- c("accrualplotdata", class(accrual_df))

  return(accrual_df)
}


#' S3 method to plot predicted recruitment from a long format 
#' dataframe of class "accrualplotdata".
#' 
#' @param data long format dataframe with columns "Week", 
#' "Arm" and "Recruitment".
#' @param plot_prefix Prefix for file name to identify plot type.
#' Defaults to `accrual_plot`.
#' @param run_time Specify a particular instance of `biomkrAccrual()`
#' execution using a date-time format `yyyy-mm-dd-hh-mm-ss`.
#' @param output_path = Directory where the output files from the 
#' `biomkrAccrual()` instance are located.
#' @param figs_path Folder where figures generated during execution
#' will be stored; defaults to the `figures` subdirectory in
#' `output_path`.
#' @param target_arm_size Number of subjects required for each treatment arm.
#' @param target_control Number of subjects required for control arm(s).
#' @param target_interim Number of subjects required for treatment arm at 
#' interim analysis.
#' @param accrual_period Number of weeks in recruitment period.
#' @param interim_period Number of weeks to recruit for interim analysis.
#' 
#' @import ggplot2
#' @importFrom grDevices palette.colors
#' @importFrom rlang .data
#' 
plot.accrualplotdata <- function(
  data,
  plot_prefix = "accrual_plot",
  run_time = NULL,
  output_path = "../biomkrAccrual_output_data/",
  figs_path = paste0(output_path, "figures/"),
  target_arm_size = NA_integer_,
  target_control = NA_integer_,
  target_interim = NA_integer_,
  accrual_period = NA_integer_,
  interim_period = NA_integer_
) {
  
  accrual_df <- data
  arm_names <- levels(accrual_df$Arm)

  linetypes <- c(
    "Interim arm" = 2, "Experimental arm" = 3, "Control arm" = 4,
    "Interim accrual" = 5, "Total accrual" = 6
  )

  hline_y <- c(target_interim, target_arm_size, target_control)
  vline_x <- c(interim_period, accrual_period)

  p <- ggplot2::ggplot(
    accrual_df, 
    ggplot2::aes(
      x = .data$Week, 
      y = .data$Recruitment, 
      group = .data$Arm, 
      color = .data$Arm
    )
  ) +
    ggplot2::geom_line(linewidth = 1) +
    # Use colourblind friendly Okabe-Ito palette
    ggplot2::scale_colour_manual(
      values = grDevices::palette.colors(length(arm_names))
    ) +
    ggplot2::geom_vline(
      xintercept = vline_x,
      linewidth = 1,
      linetype = 2:3,
      color = "grey75"
    ) +
    ggplot2::geom_hline(
      yintercept = hline_y,
      linewidth = 1,
      linetype = 4:6,
      color = "grey65"
    ) +
    ggplot2::labs(
      title = "Accrual plot"
    ) +
    theme_bma(base_size = 16)

  return(p)
}


#' Plot distributions of recruitment to arms at given time.
#' 
#' @param data Matrix with columns for each recruitment arm, 
#' including control.
#' @param target Vector of targets for recruitment. First two
#' should be those directly relevant to the subject of the graph.
#' @param target_names Vector of target names, for labelling.
#' @param target The adjust parameter from `ggplot2::geom_density`;
#' higher values mean more smoothing. Defaults to 1.
#' 
#' @importFrom stats reshape
#' @import ggplot2
#' @importFrom grDevices palette.colors
#' 
#' @export
#' 
plot.armtotals <- function(
  data,
  target,
  target_names,
  target_week,
  adjust = 1
) {
  data_df <- matrix_to_long(data)

  # Which of the accrual targets are within the dataset
  if (length(target) > 2) {
    target_indices <- 
      c(1:2, 2 + which(target[-c(1, 2)] <= max(data_df$Recruitment)))
  } else {
    target_indices <- seq_len(length(target))
  }
  target <- target[target_indices]
  target_names <- target_names[target_indices]

  p <- ggplot2::ggplot(
    data = data_df
  ) +
    ggplot2::geom_density(
      ggplot2::aes(
        x = Recruitment, group = Arm, fill = Arm, col = Arm
      ),
      alpha = 0.4, adjust = 1
    ) +
    ggplot2::scale_fill_manual(
      values = grDevices::palette.colors(length(unique(data_df$Arm))),
      aesthetics = c("color", "fill")
    ) +
    ggplot2::geom_vline(
      xintercept = target, 
      linetype = "dashed",
      linewidth = 1,
      colour = "grey65"
    ) +
    ggplot2::labs(
      y = "Probability density",
      title = target_names[1],
    ) +
    theme_bma(base_size = 16)

  p <- label_vlines(p, target, target_names)

  return(p)
}


#' Adds labels for vlines to accrual plots
#' 
#' @param p Ggplot object.
#' @param target Vector of x axis positions of vlines.
#' @param target_names Vector of target names (excluding the
#' word `target`).
#' @param size Font size (in ggplot measure) for labels; 
#' defaults to 6.
#' 
label_vlines <- function(
  p,
  target,
  target_names,
  size = 6
) {
  # Get height of y axis for this particular plot
  label_y <- round(ggplot2::layer_scales(p)$y$range$range[2], 2)

  # Get x range of plot
  xrange <- round(ggplot2::layer_scales(p)$x$range$range, 2)

  whisker <- diff(xrange) * .2

  # There's more than one target.  Move labels for the one at the range end

  # Add labels for vlines
  abline_df <- data.frame(
    x = target, 
    y = label_y * 0.9, 
    label = paste(target_names, "\ntarget")
  )

  p <- p +
    ggplot2::geom_text(
      data = abline_df,
      ggplot2::aes(x = x, y = y, label = label),
      size = size,
      family = get_base_family()
    )

  return(p)
}


#' Plot single arm accrual plot
#' 
#' Bodge, fix later
#' 
#' 
accrual_arm_plot <- function(
  data_df,
  arm_colours,
  treatment_arms,
  targets,
  plot_id,
  i
) {
  arm_names <- colnames(data_df)

  if (length(unique(data_df[, i])) == 1) {
    # BODGE - don't want to see this but need it to produce graph
    arm_col <- "white"
    alpha <- 0.0001
  } else {
    arm_col <- arm_colours[i]
    alpha <- 0.4
  }


  p <- ggplot2::ggplot(
    data = data_df
  ) +
    ggplot2::geom_density(
      ggplot2::aes(x = .data[[arm_names[i]]]),
      col = arm_col, fill = arm_col,
      alpha = alpha, adjust = 1
    ) 
  
  if (length(unique(data_df[, i])) == 1) {
    p <- p +
      ggplot2::geom_vline(
        xintercept = unique(data_df[, i]),
        linewidth = 2,
        colour = arm_colours[i],
        alpha = 0.4
      )
  }

  p <- p + 
    ggplot2::geom_vline(
      xintercept = ifelse(
        treatment_arms[i], 
        targets[1], 
        targets[2]
      ), 
      linetype = "dashed",
      linewidth = 1,
      colour = "grey75"
    ) +
    ggplot2::labs(
      x = paste(
        "No. virtual patients recruited at", 
        "target week for",
        tolower(plot_id)
      ),
      y = "Probability density",
      title = paste(plot_id, "for", arm_names[i]),
    ) +
    ggplot2::scale_x_continuous(expand = expansion(mult = 0.07)) +
    theme_bma(base_size = 16)

  p <- label_vlines(
    p, 
    target = ifelse(
      treatment_arms[i],
      targets[1], 
      targets[2]
    ),
    target_names = ifelse(
      treatment_arms[i],
      plot_id, 
      paste0(plot_id, "\ncontrol")
    )
  )

  return(p)
}