library(ggplot2)
library(dplyr)
library(ggimage)

alpha      <- 0.3
labelcolor <- "black"

# Walker image paths — copy left.png / right.png from Ch12/figures/ to Ch09/figures/ if needed,
# or simply point these at the Ch12 copies below.
img_left  <- "Ch12/figures/left.png"
img_right <- "Ch12/figures/right.png"

# ── Grid ──────────────────────────────────────────────────────────────────────
# x = ice cream scoops (0–22), y = burritos (0–11)
x_vals <- seq(0, 22, length.out = 300)
y_vals <- seq(0, 11, length.out = 300)
grid <- expand.grid(x = x_vals, y = y_vals)

# Budget constraint: y = 10 - 0.5x  →  feasible iff y ≤ 10 - 0.5x  (& y ≥ 0)
# Optimal point H: x* = 6, y* = 7  (MRS = alpha/(1-alpha) * y/x = 0.5 = price ratio)
grid <- grid %>%
  mutate(
    budget_y = pmax(10 - 0.5 * x, 0),
    valid    = y <= budget_y,
    U        = ifelse(valid & x > 0 & y > 0, x^alpha * y^(1 - alpha), NA)
  )

# Budget constraint line for geom_line
bc_line <- data.frame(
  x = seq(0, 20, length.out = 400),
  y = 10 - 0.5 * seq(0, 20, length.out = 400)
)

# ── Plot ──────────────────────────────────────────────────────────────────────
textbook_plot <- function() {
  ggplot(grid, aes(x = x, y = y)) +

    # Heatmap — only feasible region; gray panel fills the infeasible area
    geom_tile(data = grid %>% filter(!is.na(U)), aes(fill = U)) +
    scale_fill_viridis_c(
      option = "C",
      name   = "Least preferred                      Most preferred",
      guide  = guide_colorbar(
        title.position = "top",
        barwidth       = 20,
        ticks          = FALSE,
        label          = FALSE
      )
    ) +

    # Budget constraint line
    geom_line(data = bc_line, aes(x = x, y = y),
              color = "black", size = 1.2, inherit.aes = FALSE) +

    # ── Point H (6, 7): utility-maximising point ──────────────────────────
    geom_segment(aes(x = 0, xend = 6, y = 7, yend = 7),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_segment(aes(x = 6, xend = 6, y = 0, yend = 7),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_point(aes(x = 6, y = 7),  shape = 16, size = 3, color = "black") +
    annotate("text", x = 6.3, y = 7.15, label = "H", size = 5, fontface = "bold") +

    # ── Point K (2, 9): upper-left ──────────────────────────────────────
    geom_segment(aes(x = 0, xend = 2, y = 9, yend = 9),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_segment(aes(x = 2, xend = 2, y = 0, yend = 9),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_point(aes(x = 2, y = 9),  shape = 16, size = 3, color = "black") +
    annotate("text", x = 1.7, y = 8.85, label = "K", size = 5, fontface = "bold") +

    # ── Point L (4, 8): between K and H ─────────────────────────────────
    geom_segment(aes(x = 0, xend = 4, y = 8, yend = 8),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_segment(aes(x = 4, xend = 4, y = 0, yend = 8),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_point(aes(x = 4, y = 8),  shape = 16, size = 3, color = "black") +
    annotate("text", x = 3.7, y = 7.85, label = "L", size = 5, fontface = "bold") +

    # ── Point M (10, 5): between H and J ────────────────────────────────
    geom_segment(aes(x = 0, xend = 10, y = 5, yend = 5),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_segment(aes(x = 10, xend = 10, y = 0, yend = 5),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_point(aes(x = 10, y = 5), shape = 16, size = 3, color = "black") +
    annotate("text", x = 9.7, y = 4.85, label = "M", size = 5, fontface = "bold") +

    # ── Point J (14, 3): lower-right ─────────────────────────────────────
    geom_segment(aes(x = 0, xend = 14, y = 3, yend = 3),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_segment(aes(x = 14, xend = 14, y = 0, yend = 3),
                 color = "gray40", linetype = "dashed", size = 0.5) +
    geom_point(aes(x = 14, y = 3), shape = 16, size = 3, color = "black") +
    annotate("text", x = 13.7, y = 2.85, label = "J", size = 5, fontface = "bold") +

    # ── Walkers near J and M (moving left-up toward H) ───────────────────
    # Placed just above the budget line (infeasible / gray zone)
    geom_image(data = data.frame(x = 14,   y = 3.7, image = img_left),
               aes(x = x, y = y, image = image), size = 0.09, inherit.aes = FALSE) +
    geom_image(data = data.frame(x = 10.8, y = 5.5, image = img_left),
               aes(x = x, y = y, image = image), size = 0.09, inherit.aes = FALSE) +

    # Arrows: J → M → H  (slope -0.5 throughout)
    geom_segment(aes(x = 13.5, y = 3.25, xend = 11.5, yend = 4.25),
                 arrow = arrow(length = unit(0.3, "cm"), type = "closed"),
                 color = "black", linewidth = 1) +
    geom_segment(aes(x = 10.7, y = 4.65, xend = 7.2,  yend = 6.4),
                 arrow = arrow(length = unit(0.3, "cm"), type = "closed"),
                 color = "black", linewidth = 1) +

    # ── Walkers near K and L (moving right-down toward H) ────────────────
    geom_image(data = data.frame(x = 2, y = 9.65, image = img_right),
               aes(x = x, y = y, image = image), size = 0.09, inherit.aes = FALSE) +
    geom_image(data = data.frame(x = 4, y = 8.65, image = img_right),
               aes(x = x, y = y, image = image), size = 0.09, inherit.aes = FALSE) +

    # Arrows: K → L → H  (slope -0.5 throughout)
    geom_segment(aes(x = 2.5,  y = 8.75, xend = 3.7, yend = 8.15),
                 arrow = arrow(length = unit(0.3, "cm"), type = "closed"),
                 color = "black", linewidth = 1) +
    geom_segment(aes(x = 4.3,  y = 7.85, xend = 5.65, yend = 7.18),
                 arrow = arrow(length = unit(0.3, "cm"), type = "closed"),
                 color = "black", linewidth = 1) +

    # ── Axis labels ───────────────────────────────────────────────────────
    labs(x = "Ice Cream Scoops (x)", y = "Burritos (y)", title = "") +

    scale_x_continuous(
      breaks = seq(0, 22, by = 2),
      limits = c(0, 22),
      expand = c(0, 0)
    ) +
    scale_y_continuous(
      breaks = seq(0, 11, by = 1),
      limits = c(0, 11),
      expand = c(0, 0)
    ) +

    # ── Theme ─────────────────────────────────────────────────────────────
    theme_minimal() +
    theme(
      panel.background = element_rect(fill = "gray80", color = NA),
      plot.background  = element_rect(fill = "white",  color = NA),
      plot.title       = element_text(hjust = 0.5, size = 18, face = "bold"),
      axis.title       = element_text(size = 14, face = "bold"),
      axis.text        = element_text(size = 12),
      legend.position  = "top",
      legend.title     = element_text(size = 14, face = "bold"),
      legend.margin    = margin(b = 10),
      plot.margin      = margin(10, 20, 10, 10),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      axis.line        = element_line(color = "black", linewidth = 0.6),
      axis.ticks       = element_line(color = "black")
    ) +

    # ── "Budget Constraint" label (bottom-right, near line) ──────────────
    annotate("text", x = 18.5, y = 2,
             label    = "Budget\nConstraint",
             color    = labelcolor,
             size     = 5,
             fontface = "bold") +

    # ── "Getting warmer" labels (in infeasible gray zone, above the line) ──
    annotate("text", x = 3.8, y = 9.8,
             label    = "Getting warmer\n(more utility)",
             color    = labelcolor,
             size     = 4,
             fontface = "bold") +
    annotate("text", x = 12.5, y = 5.6,
             label    = "Getting warmer\n(more utility)",
             color    = labelcolor,
             size     = 4,
             fontface = "bold")
}

# ── Save ──────────────────────────────────────────────────────────────────────
ggsave(
  "Ch09/figures/heatmap_bc_navigation.png",
  plot   = textbook_plot(),
  width  = 10,
  height = 7,
  dpi    = 300,
  bg     = "white"
)

ggsave(
  "Ch09/figures/heatmap_bc_navigation.pdf",
  plot   = textbook_plot(),
  width  = 10,
  height = 7,
  bg     = "white"
)