library(ggplot2)
library(dplyr)

# 1. Set your Cobb–Douglas exponent
alpha <- 0.3
curvecolor <- "black"
pointcolor <- "black"
labelcolor <- "black"

# 2. Create a high-res grid
grid <- expand.grid(
  x = seq(0.1, 22, length.out = 300),
  y = seq(0.1, 11, length.out = 300)
)
grid$U <- grid$x^alpha * grid$y^(1 - alpha)

# Budget constraint: 2x + 3y = 30
# Extending from y=10 to x=15 (where y=0)
budget_df <- data.frame(
  x = seq(0, 20, 0.1),
  y = (10 - 0.5 * seq(0, 20, 0.1))
)
budget_df <- budget_df[budget_df$y >= 0, ]

# Points H, I, J, K, L
points_df <- data.frame(
  x = c(6, 6, 12, 12),  # K is on budget constraint at (12, 2)
  y = c(4, 7, 4, 7),     # L is beyond constraint
  label = c("H", "I", "J", "K")
)

# Textbook-ready plot with fixed issues
textbook_plot <- function() {
  ggplot(grid, aes(x = x, y = y, fill = U)) +
    # Background heatmap
    geom_tile() +
    scale_fill_viridis_c(
      option = "C",
      name = "Least preferred                      Most preferred",
      guide = guide_colorbar(
        title.position = "top",
        barwidth = 20,
        ticks = FALSE,
        label = FALSE
      )
    ) +

    # Axis formatting with expanded limits to prevent cropping
    scale_x_continuous(
      breaks = seq(0, 22, by = 4),
      limits = c(0, 22),
      expand = c(0.01, 0.5)  # Increased right margin
    ) +
    scale_y_continuous(
      breaks = seq(0, 11, by = 2),
      limits = c(0, 11),
      expand = c(0.01, 0.2)  # Increased top margin
    ) +

    # Add budget constraint line
    geom_line(
      data = budget_df,
      aes(x = x, y = y),
      inherit.aes = FALSE,
      color = curvecolor,
      size = 1.2
    ) +

    # Add points
    geom_point(
      data = points_df,
      aes(x = x, y = y),
      inherit.aes = FALSE,
      color = pointcolor,
      size = 3,
      stroke = 1
    ) +

    # Add labels for points
    geom_text(
      data = points_df,
      aes(x = x, y = y, label = label),
      inherit.aes = FALSE,
      color = labelcolor,
      size = 5,
      nudge_x = 0.6,
      fontface = "bold"
    ) +

    # Labels and title
    labs(
      x = "Ice Cream Scoops (x)",
      y = "Burritos (y)",
      title = ""
    ) +

    # Theme customization for textbook appearance
    theme_minimal() +
    theme(
      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)  # Increased right margin
    ) +

    # Moved budget constraint label to avoid overlap
    annotate(
      "text",
      x = 6, #Moved to the left side
      y = 8.5,    # Moved higher up
      label = "Budget Constraint",
      color = labelcolor,
      size = 4.5,
      fontface = "bold"
    )
}

# Save high-resolution plot with expanded canvas
ggsave(
  "Ch09/figures/heatmap_bc.png",
  plot = textbook_plot(),
  width = 10,  # Wider canvas
  height = 7,  # Taller canvas
  dpi = 300,
  bg = "white"
)

# For PDF output (vector graphics, better for publishing)
ggsave(
  "Ch09/figures/heatmap_bc.pdf",
  plot = textbook_plot(),
  width = 10,  # Wider canvas
  height = 7,  # Taller canvas
  bg = "white"
)