## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4
)

## -----------------------------------------------------------------------------
library(PiChange)
library(ggplot2)

data("google_mobility_workplaces")

run_region <- function(region, event_dates) {
  region_data <- google_mobility_workplaces[
    google_mobility_workplaces$region == region,
  ]
  y <- region_data$workplaces_percent_change_from_baseline
  penalty <- construct_penalty(
    time = region_data$date,
    centers = as.Date(event_dates),
    width = 30,
    method = "mbic",
    family = "normal"
  )
  fit <- pi_change(y, penalty, min_seg_len = 30)
  list(
    data = region_data,
    fit = fit,
    centers = penalty$prior$center_indices,
    change_indices = changepoints(fit)
  )
}

us_events <- c("2020-03-15", "2020-05-01", "2020-12-11",
               "2021-03-11", "2022-08-16")
uk_events <- c("2020-03-23", "2020-06-08", "2020-10-31",
               "2020-12-02", "2020-12-08", "2020-12-26",
               "2021-03-29", "2022-02-24")

fits <- list(
  "United States" = run_region("United States", us_events),
  "New York City" = run_region("New York City", us_events),
  "United Kingdom" = run_region("United Kingdom", uk_events),
  "London" = run_region("London", uk_events)
)

data.frame(
  region = rep(names(fits), lengths(lapply(fits, `[[`, "change_indices"))),
  date = do.call(c, lapply(fits, function(x) x$data$date[x$change_indices]))
)

## ----fig.height=6-------------------------------------------------------------
plot(
  fits[["United States"]]$fit,
  ylab = "Workplace mobility percent change",
  title = "United States workplace mobility"
)

## -----------------------------------------------------------------------------
plot_data <- do.call(rbind, lapply(names(fits), function(region) {
  x <- fits[[region]]
  data.frame(
    date = x$data$date,
    region = region,
    workplaces_percent_change_from_baseline =
      x$data$workplaces_percent_change_from_baseline,
    segment_mean = fitted(x$fit),
    pi_change = seq_len(nrow(x$data)) %in% x$change_indices,
    prior_center = seq_len(nrow(x$data)) %in% x$centers
  )
}))

ggplot(plot_data, aes(x = date, y = workplaces_percent_change_from_baseline)) +
  geom_line(
    aes(color = "Observed series", linetype = "Observed series"),
    linewidth = 0.4
  ) +
  geom_line(
    aes(y = segment_mean, color = "Segment mean", linetype = "Segment mean"),
    linewidth = 0.4
  ) +
  geom_vline(
    data = plot_data[plot_data$pi_change, ],
    aes(
      xintercept = date,
      color = "Detected change",
      linetype = "Detected change"
    ),
    linewidth = 0.8
  ) +
  geom_vline(
    data = plot_data[plot_data$prior_center, ],
    aes(
      xintercept = date,
      color = "Prior center",
      linetype = "Prior center"
    ),
    linewidth = 0.8
  ) +
  facet_wrap(~ region, ncol = 2) +
  scale_color_manual(
    values = c(
      "Observed series" = "#76BDE8",
      "Segment mean" = "#F06292",
      "Prior center" = "#D62728",
      "Detected change" = "#00C853"
    )
  ) +
  scale_linetype_manual(
    values = c(
      "Observed series" = "solid",
      "Segment mean" = "solid",
      "Prior center" = "dashed",
      "Detected change" = "solid"
    )
  ) +
  scale_x_date(date_breaks = "6 months", date_labels = "%b %Y") +
  labs(
    x = "Date",
    y = "Workplace mobility percent change",
    title = "Workplace mobility with PI-Change estimates",
    color = NULL,
    linetype = NULL
  ) +
  theme_minimal(base_size = 11) +
  theme(
    panel.grid.minor = element_blank(),
    panel.grid.major.x = element_blank(),
    legend.position = "top",
    legend.justification = "left",
    axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1)
  )