# =====================================================================
#  Recovering Trust in AI vs. Human Experts After Short-Term and
#  Long-Term Shocks  --  Discussion Figures 1-4
#
#  What the figures establish
#  --------------------------
#  Fig 1  The SAME shock produces DIFFERENT observed series depending on
#         the data-generating coefficients.  AR (phi) is inertia, MA
#         (theta) is rebound; when they are equal they cancel and the
#         process collapses to a pure random walk, ARIMA(0,1,0).
#  Fig 2  What we actually face: an ARIMA(0,1,0)-type permanent shock
#         arriving on top of a pre-existing ARIMA(0,1,1) structure.  The
#         advisor keeps its habitual updating rule, so the forecast line
#         lags behind the observed level for several periods.
#  Fig 3  That lag, isolated: the forecast-error path.
#  Fig 4  Where the manipulation is identified at all, over the
#         advisor's (phi0, theta0) grid.
#
#  Author: AIForA Lab (Dae Keun Park)
#  Run:    Rscript trust_recovery_figures.R
#  Output: figures/fig1..fig4  as PNG (paper + web themes) and PDF
# =====================================================================

OUT <- "figures"
dir.create(OUT, showWarnings = FALSE)

## ---- core algebra ---------------------------------------------------

# Level impulse response of ARIMA(1,1,1) to a single innovation eps_80 = M.
#   (1 - phi B)(1 - B) y_t = (1 - theta B) eps_t
#   y_{80+k}/M = 1 + (phi - theta) (1 - phi^k)/(1 - phi)
#   long-run multiplier = (1 - theta)/(1 - phi)
irf_level <- function(phi, theta, K = 13) {
  k <- 0:K
  s <- if (phi == 1) k else (1 - phi^k) / (1 - phi)
  c(1, (1 + (phi - theta) * s)[-1])
}

longrun <- function(phi, theta) (1 - theta) / (1 - phi)

# One-step forecast error when a permanent (random-walk type) shock of
# size M hits and the advisor keeps a habitual ARIMA(1,1,1) rule:
#   e_t = w_t + theta0 e_{t-1}
#   e_80 = M ,  e_{80+k} = M (theta0 - phi0) theta0^(k-1)
fc_error <- function(phi0, theta0, K = 13)
  c(1, (theta0 - phi0) * theta0^(0:(K - 1)))

# Post-shock periods in which |e| is still at least 20% of the shock.
# NOTE: 0.2 is a stated threshold, not a point of "arrival" -- exponential
# smoothing approaches the new level asymptotically and never reaches it.
# theta0 = 0.7 gives 4 periods at this threshold (2 at 50%, 7 at 10%, 13 at 1%).
n_persist <- function(phi0, theta0)
  sum(abs((theta0 - phi0) * theta0^(0:19)) >= 0.2)

## ---- coefficient cases shown in every panel -------------------------

CASES <- data.frame(
  phi   = c(0.0, 0.0, 0.3, 0.0, 0.5),
  theta = c(0.7, 0.3, 0.7, 0.0, 0.3)
)
# Figure 1 varies the DATA-GENERATING process, so its coefficients are
# (phi, theta).  Figures 2-3 hold the data fixed and vary the ADVISOR's
# habitual updating rule, whose coefficients are (phi0, theta0).  The two
# are different objects, so they get different symbols and legend titles.
# Plain meaning first, coefficients in parentheses -- the reader should be
# able to read the legend without already knowing what phi and theta do.
CASES$note1 <- c("keeps 30% of the shock",
                 "keeps 70% of the shock",
                 "keeps 43% of the shock",
                 "keeps all of it - random walk",
                 "overshoots to 140%")
CASES$note2 <- c("lags behind 4 periods",
                 "catches up after 1 period",
                 "lags behind 2 periods",
                 "catches up at once",
                 "overshoots, error flips sign")

CASES$lab_dgp <- sprintf("%s   (phi %.1f, theta %.1f)",
                         CASES$note1, CASES$phi, CASES$theta)
CASES$lab_adv <- sprintf("%s   (phi0 %.1f, theta0 %.1f)",
                         CASES$note2, CASES$phi, CASES$theta)

T0 <- 76; T1 <- 93; SHOCK <- 80
tt <- T0:T1
kk <- tt - SHOCK

## ---- themes ---------------------------------------------------------

theme_paper <- list(bg = "#ffffff", fg = "#111111", mut = "#666666", tan = "#8a5a1f",
                    grid = "#e2e2e2", axis = "#999999",
                    col = c("#2a78d6", "#c2620f", "#12855c",
                            "#6b7280", "#c2185b"))
theme_web   <- list(bg = "#12213a", fg = "#eaf1fb", mut = "#9db2d4", tan = "#e0a96d",
                    grid = "#25406b", axis = "#3d5c8f",
                    col = c("#7fb3ff", "#e0a96d", "#6fd6b0",
                            "#b9c6dd", "#f08fb0"))

## ---- shared line-panel drawing --------------------------------------

line_panel <- function(th, ylim, yticks, valfun, ylab, notes, labs,
                       leg_title, show_obs = FALSE, lwd1 = 3.8) {
  par(bg = th$bg, fg = th$fg, col.axis = th$mut, col.lab = th$mut,
      mar = c(11.5, 5.4, 1.4, 2.2), xpd = FALSE)
  plot(NA, xlim = c(T0, T1 + 0.2), ylim = ylim, xlab = "", ylab = "",
       axes = FALSE)
  abline(h = yticks, col = th$grid, lwd = 1)
  if (0 >= ylim[1] && 0 <= ylim[2]) abline(h = 0, col = th$axis, lwd = 1.4)
  abline(v = SHOCK, col = th$axis, lwd = 1, lty = 3)
  axis(1, at = seq(T0, T1, 2), col = th$axis, col.ticks = th$axis, cex.axis = 1.0)
  axis(2, at = yticks, las = 1, col = th$axis, col.ticks = th$axis, cex.axis = 1.0)
  mtext("period t", side = 1, line = 3.0, col = th$mut, cex = 1.05)
  mtext(ylab, side = 2, line = 3.7, col = th$mut, cex = 1.05)
  text(SHOCK + .25, ylim[2] - diff(ylim) * .04, "shock", adj = 0,
       col = th$mut, cex = .95)

  if (show_obs) {
    obs <- ifelse(kk < 0, 0, 1)
    lines(tt, obs, col = th$fg, lwd = 2.6, lty = 2, type = "s")
  }

  for (i in seq_len(nrow(CASES))) {
    v <- sapply(kk, function(k) valfun(CASES$phi[i], CASES$theta[i], k))
    lines(tt, v, col = th$col[i], lwd = if (i == 1) lwd1 else 2.6)
    points(tt[kk >= 0], v[kk >= 0], pch = 21, bg = th$col[i],
           col = th$bg, cex = 1.0, lwd = 1.6)
  }

  # Legend lives in the reserved right margin, so it can never collide
  # with the curves -- which converge at the right edge in figs 2 and 3.
  par(xpd = NA)
  keys <- labs
  cols <- th$col
  lwds <- c(lwd1, rep(2.6, nrow(CASES) - 1))
  ltys <- rep(1, nrow(CASES))
  ttl  <- leg_title
  if (show_obs) {
    keys <- c("the actual data - jumps once, stays there", keys)
    cols <- c(th$fg, cols); lwds <- c(2.6, lwds); ltys <- c(2, ltys)
  }
  legend(x = T0 - 0.8, y = ylim[1] - diff(ylim) * 0.28, legend = keys, col = cols,
         lwd = lwds + 0.8, lty = ltys, bty = "n", cex = 0.95, seg.len = 2.6,
         ncol = 2, x.intersp = 0.9, y.intersp = 1.5, text.col = th$fg,
         title = ttl, title.col = th$tan, title.adj = 0, title.cex = 1.05)
  par(xpd = FALSE)
}

## ---- figure 4: identification map -----------------------------------

grid_panel <- function(th) {
  g   <- seq(0, 0.9, length.out = 46)
  z   <- outer(g, g, Vectorize(n_persist))          # rows = phi0, cols = theta0
  ramp <- colorRampPalette(if (identical(th$bg, "#ffffff"))
    c("#f2f5fa", "#cde2fb", "#9ec5f4", "#5598e7", "#2a78d6", "#1c5cab", "#0d366b")
    else c("#16243c", "#1d3a63", "#25538e", "#3f8ae0", "#6aa8ec", "#9cc7f4", "#cfe3fb"))
  par(bg = th$bg, fg = th$fg, col.axis = th$mut, mar = c(10.5, 5.4, 1.4, 2.2))
  image(g, g, z, col = ramp(max(z) + 1), axes = FALSE, xlab = "", ylab = "",
        zlim = c(0, max(z)))
  axis(1, at = seq(0, .9, .3), col = th$axis, col.ticks = th$axis, cex.axis = 1.0)
  axis(2, at = seq(0, .9, .3), las = 1, col = th$axis, col.ticks = th$axis, cex.axis = 1.0)
  mtext(expression("advisor AR coefficient " * phi[0]), 1, line = 3.1, col = th$mut, cex = 1.05)
  mtext(expression("advisor MA coefficient " * theta[0]), 2, line = 3.8, col = th$mut, cex = 1.05)
  abline(0, 1, col = th$fg, lwd = 1.5, lty = 2)
  points(0, 0.7, pch = 1, cex = 2.4, col = th$fg, lwd = 2)
  text(0.05, 0.735, "this study: advisor theta0 = 0.7", adj = 0, col = th$fg, cex = 1.0)
  text(0.05, 0.672, "error stays above 20% of the shock for 4 periods", adj = 0, col = th$fg, cex = .88)
  text(0.57, 0.83, expression(theta[0] > phi[0]), adj = 0, col = th$fg, cex = 1.2)
  text(0.57, 0.77, "manipulation identified", adj = 0, col = th$fg, cex = .95)
  text(0.23, 0.13, expression(theta[0] <= phi[0]), adj = 0, col = th$fg, cex = 1.2)
  text(0.23, 0.07, "vanishes or reverses sign", adj = 0, col = th$fg, cex = .95)

  # discrete colour key in the reserved right margin
  par(xpd = NA)
  cols <- ramp(max(z) + 1)
  brk  <- rev(seq(0, max(z), length.out = min(5, max(z) + 1)))
  brk  <- unique(round(brk))
  legend(x = -0.07, y = -0.20,
         legend = paste(brk, ifelse(brk == 1, "period", "periods")),
         fill = cols[brk + 1], border = NA, bty = "n", cex = 0.95,
         ncol = 5, x.intersp = 0.8, text.col = th$fg,
         title = "periods until the error falls below 20% of the shock (no correction)",
         title.col = th$tan, title.adj = 0, title.cex = 1.05)
  par(xpd = FALSE)
}

## ---- render ---------------------------------------------------------

figures <- list(
  fig1 = function(th) line_panel(th, c(-0.05, 1.55), seq(0, 1.5, .25),
           function(p, q, k) if (k < 0) 0 else irf_level(p, q, 20)[k + 1],
           "observed level  y / M", CASES$note1, CASES$lab_dgp,
           "how the DATA is generated", show_obs = FALSE),
  fig2 = function(th) line_panel(th, c(-0.05, 1.35), seq(0, 1.25, .25),
           function(p, q, k) if (k < 0) 0 else 1 - fc_error(p, q, 20)[k + 1],
           "level  y / M", CASES$note2, CASES$lab_adv,
           "what the ADVISOR forecasts", show_obs = TRUE),
  fig3 = function(th) line_panel(th, c(-0.32, 1.05), seq(-0.25, 1, .25),
           function(p, q, k) if (k < 0) 0 else fc_error(p, q, 20)[k + 1],
           "forecast error  e / M", CASES$note2, CASES$lab_adv,
           "what the ADVISOR forecasts", show_obs = FALSE),
  fig4 = function(th) grid_panel(th)
)

for (nm in names(figures)) {
  png(file.path(OUT, paste0(nm, "_paper.png")), width = 2200, height = 1300,
      res = 180, pointsize = 15, bg = theme_paper$bg)
  figures[[nm]](theme_paper); dev.off()

  png(file.path(OUT, paste0(nm, "_web.png")), width = 2200, height = 1300,
      res = 180, pointsize = 15, bg = theme_web$bg)
  figures[[nm]](theme_web); dev.off()

  pdf(file.path(OUT, paste0(nm, ".pdf")), width = 12.2, height = 7.2, pointsize = 15)
  figures[[nm]](theme_paper); dev.off()
}

## ---- numbers quoted in the discussion -------------------------------

cat("\n--- Figure 1: long-run multiplier (1-theta)/(1-phi) ---\n")
print(data.frame(CASES[, c("phi", "theta")],
                 retained = round(mapply(longrun, CASES$phi, CASES$theta), 4)))

cat("\n--- Figure 3: forecast-error path e/M, k = 0..5 ---\n")
e <- t(mapply(function(p, q) round(fc_error(p, q, 5), 4), CASES$phi, CASES$theta))
dimnames(e) <- list(CASES$lab, paste0("k=", 0:5))
print(e)

cat("\n--- Figure 4: periods with |e| >= 0.2M ---\n")
print(data.frame(CASES[, c("phi", "theta")],
                 periods = mapply(n_persist, CASES$phi, CASES$theta)))

cat("\n--- cancellation check: phi = theta gives a pure random walk ---\n")
cat("phi=theta=0.7 :", round(irf_level(0.7, 0.7, 5), 4), "\n")
cat("phi=theta=0.0 :", round(irf_level(0.0, 0.0, 5), 4), "\n")
cat("identical     :", isTRUE(all.equal(irf_level(0.7, 0.7), irf_level(0, 0))), "\n")

cat("\nWritten to", normalizePath(OUT), "\n")