# Visualisation helpers for wid_df objects

.PAL <- c("#58a6ff", "#3fb950", "#d2a8ff", "#ffa657", "#ff7b72",
          "#79c0ff", "#56d364", "#e3b341")

.theme_wid <- function() theme(
  plot.background   = element_rect(fill = "#0d1117", color = NA),
  panel.background  = element_rect(fill = "#0d1117", color = NA),
  panel.grid.major  = element_line(color = "#21262d"),
  panel.grid.minor  = element_blank(),
  axis.text         = element_text(color = "#8b949e"),
  axis.title        = element_text(color = "#e6edf3"),
  plot.title        = element_text(color = "#e6edf3", face = "bold"),
  plot.subtitle     = element_text(color = "#8b949e"),
  legend.background = element_rect(fill = "#161b22", color = NA),
  legend.text       = element_text(color = "#e6edf3"),
  legend.title      = element_text(color = "#8b949e"),
  strip.background  = element_rect(fill = "#161b22"),
  strip.text        = element_text(color = "#e6edf3")
)

.prep <- function(data, variable, country, country_labels) {
  .check(data)
  variable <- .resolve_var(data, variable)
  d <- data[data$variable == variable, ]
  if (!is.null(country)) d <- d[d$country %in% country, ]
  if (!nrow(d)) stop("No data after filtering.", call. = FALSE)
  if (!is.null(country_labels)) d$country <- .recode(d$country, country_labels)
  list(d = d, v = variable)
}

#' Time-series line chart for WID data
#'
#' @param data           A \code{wid_df}.
#' @param variable       Variable code to plot.
#' @param country        Country codes to include.
#' @param country_labels Named character vector for renaming countries.
#' @param facet          Logical. Facet by country. Default \code{FALSE}.
#' @return A \code{ggplot} object.
#' @examples
#' \donttest{
#'   d <- data.frame(
#'     country = rep(c("US", "FR"), each = 3L),
#'     year    = rep(c("2018", "2019", "2020"), 2L),
#'     variable = "sptinc992j", percentile = "p99p100",
#'     value = c(0.20, 0.19, 0.19, 0.11, 0.11, 0.12),
#'     age = "992", pop = "j", stringsAsFactors = FALSE
#'   )
#'   class(d) <- c("wid_df", "data.frame")
#'   wid_plot_timeseries(d)
#' }
#' @export
wid_plot_timeseries <- function(data, variable = NULL, country = NULL,
                                country_labels = NULL, facet = FALSE) {
  p <- .prep(data, variable, country, country_labels)
  g <- ggplot(p$d, aes(year, value, color = country)) +
    geom_line(linewidth = 0.8) +
    geom_point(size = 1.5) +
    scale_color_manual(values = .PAL) +
    labs(title = .auto_title(p$v, unique(p$d$country)),
         x = "Year", y = "Value", color = NULL) +
    .theme_wid()
  if (facet) g + facet_wrap(~country) else g
}

#' Cross-country comparison bar/point chart
#'
#' @param data           A \code{wid_df}.
#' @param year           Integer year (defaults to most recent common year).
#' @param variable       Variable code to plot.
#' @param country        Country codes to include.
#' @param country_labels Named character vector for renaming countries.
#' @param type           \code{"bar"} (default) or \code{"point"}.
#' @return A \code{ggplot} object.
#' @examples
#' \donttest{
#'   d <- data.frame(
#'     country = rep(c("US", "FR"), each = 3L),
#'     year    = rep(c("2018", "2019", "2020"), 2L),
#'     variable = "sptinc992j", percentile = "p99p100",
#'     value = c(0.20, 0.19, 0.19, 0.11, 0.11, 0.12),
#'     age = "992", pop = "j", stringsAsFactors = FALSE
#'   )
#'   class(d) <- c("wid_df", "data.frame")
#'   wid_plot_compare(d, year = 2020)
#' }
#' @export
wid_plot_compare <- function(data, year = NULL, variable = NULL, country = NULL,
                             country_labels = NULL, type = "bar") {
  p   <- .prep(data, variable, country, country_labels)
  yr  <- year %||% .common_year(p$d)
  d   <- p$d[p$d$year == yr, ]
  if (!nrow(d)) stop("No data for year ", yr, ".", call. = FALSE)
  agg <- tapply(d$value, d$country, mean, na.rm = TRUE)
  d$country <- factor(d$country, levels = names(sort(agg, decreasing = TRUE)))
  g <- ggplot(d, aes(country, value)) +
    labs(title = paste0(.auto_title(p$v, levels(d$country)), " (", yr, ")"),
         x = NULL, y = "Value") +
    .theme_wid() +
    theme(legend.position = "none")
  if (type == "bar")
    g + geom_col(aes(fill = country)) + scale_fill_manual(values = .PAL)
  else
    g + geom_point(size = 3, color = .PAL[1L])
}

#' Lorenz curve
#'
#' @param data           A \code{wid_df} with share (\code{"s"}) series.
#' @param variable       Variable code to plot.
#' @param country        Country codes to include.
#' @param country_labels Named character vector for renaming countries.
#' @return A \code{ggplot} object.
#' @examples
#' \donttest{
#'   d <- data.frame(
#'     country = rep("US", 3L), year = rep("2020", 3L),
#'     variable = "sptinc992j",
#'     percentile = c("p0p50", "p50p90", "p90p100"),
#'     value = c(0.20, 0.40, 0.40), age = "992", pop = "j",
#'     stringsAsFactors = FALSE
#'   )
#'   class(d) <- c("wid_df", "data.frame")
#'   wid_plot_lorenz(d)
#' }
#' @export
wid_plot_lorenz <- function(data, variable = NULL, country = NULL,
                            country_labels = NULL) {
  .check(data, series_type = "s")
  p <- .prep(data, variable, country, country_labels)
  d <- p$d
  
  bounds <- lapply(d$percentile, .parse_perc)
  ok     <- !vapply(bounds, is.null, logical(1L))
  if (any(!ok))
    warning(sum(!ok), " row(s) dropped: unparseable percentile codes.", call. = FALSE)
  d     <- d[ok, ]
  d$hi  <- vapply(bounds[ok], `[[`, numeric(1L), "hi")
  d     <- do.call(rbind, lapply(split(d, d$country), function(g) {
    g <- g[order(g$hi), ]
    g$cum_share <- cumsum(g$value)
    g
  }))
  
  ggplot(d, aes(hi, cum_share, color = country)) +
    geom_abline(slope = 1, intercept = 0, linetype = "dashed",
                color = "#8b949e", linewidth = 0.5) +
    geom_line(linewidth = 0.9) +
    scale_color_manual(values = .PAL) +
    scale_x_continuous(labels = percent_format()) +
    scale_y_continuous(labels = percent_format()) +
    labs(title = "Lorenz Curve", x = "Population share",
         y = "Income/wealth share", color = NULL) +
    .theme_wid()
}