# ------------------------------------------------------------
# Replication of "The Model and its properties" 
# 
# ------------------------------------------------------------

rm(list = ls())

library(dplyr)
library(tidyr)
library(ggplot2)

# ---- Model simulator ----
# Endogenous: p, w, p_e, pi_star
# Inflation:  pi_t = p_t - p_{t-1}
#
# Equations (from slides):
# (1) w_t - w_{t-1} = (p^e_t - p_{t-1}) + α(p_{t-1} - p^e_{t-1}) + β(x_t - α x_{t-1}) + z^w_t
# (2) p_t - p_{t-1} = (w_t - w_{t-1}) + (z^p_t - z^p_{t-1})
# (3) p^e_t - p_{t-1} = δ pi*_t + (1-δ)(p_{t-1} - p_{t-2})
# (4) pi*_t = γ pi*_{t-1} + (1-γ)(p_{t-1} - p_{t-2})

simulate_model <- function(T = 16,
                           alpha, beta, delta, gamma,
                           zp_path, x_path,
                           zw_path = rep(0, length(zp_path))) {
  
  stopifnot(length(zp_path) == T,
            length(x_path)  == T,
            length(zw_path) == T)
  
  # Allocate (include t=0 and t=-1 lags via indices)
  p       <- rep(0, T + 2)  # p[1]=t=-1, p[2]=t=0, p[t+2]=t
  w       <- rep(0, T + 1)  # w[1]=t=0
  p_e     <- rep(0, T + 1)  # p_e[1]=t=0
  pi_star <- rep(0, T + 1)  # pi_star[1]=t=0
  zp      <- c(0, zp_path)   # zp[1]=t=0, zp[t+1]=t
  x       <- c(0, x_path)    # x[1]=t=0
  
  pi <- rep(0, T)  # inflation pi_t for t=1..T
  
  for (t in 1:T) {
    # Map indices:
    # p_{t-1} = p[t+1], p_{t-2} = p[t]
    p_tm1 <- p[t + 1]
    p_tm2 <- p[t]
    pi_tm1 <- p_tm1 - p_tm2
    
    # (4) pi*_t
    pi_star[t + 1] <- gamma * pi_star[t] + (1 - gamma) * pi_tm1
    
    # (3) p^e_t
    p_e[t + 1] <- p_tm1 + delta * pi_star[t + 1] + (1 - delta) * (p_tm1 - p_tm2)
    
    # (1) wage growth
    dw_t <- (p_e[t + 1] - p_tm1) +
      alpha * (p_tm1 - p_e[t]) +
      beta  * (x[t + 1] - alpha * x[t]) +
      zw_path[t]
    
    # update w
    w[t + 1] <- w[t] + dw_t
    
    # (2) inflation and update p
    dzp_t <- zp[t + 1] - zp[t]
    pi[t] <- dw_t + dzp_t
    p[t + 2] <- p_tm1 + pi[t]
  }
  
  tibble(
    quarter = 1:T,
    inflation = pi
  )
}

# ---- Shock designs (match the slide experiments) ----
# 1) Price shock: a one-time permanent level shift in z^p at t=1:
#    z^p_0 = 0, z^p_1 = shock, and stays there => (z^p_t - z^p_{t-1}) is nonzero only at t=1.
make_price_shock <- function(T, shock_size = 1) {
  zp <- rep(shock_size, T)
  zp[1] <- shock_size  # jump at t=1, then permanent
  zp
}

# 2) Permanent increase in labor market tightness x: step at t=1 and stays
make_perm_x <- function(T, shock_size = 1) rep(shock_size, T)

# ---- Parameter scenarios from slide ----
# Weak feedback:  α = 0.2, δ = 0.9, γ = 0.95
# Strong feedback: α = 0.6, δ = 0.7, γ = 0.9
# NOTE: β is not shown in the screenshot; set it here (and tweak if you want exact levels).
beta_default <- 0.5

par_weak   <- list(alpha = 0.2, beta = beta_default, delta = 0.9, gamma = 0.95)
par_strong <- list(alpha = 0.6, beta = beta_default, delta = 0.7, gamma = 0.90)

T <- 16

# ---- IRF 1: Inflation response to price shock ----
zp1 <- make_price_shock(T, shock_size = 1)
x0  <- rep(0, T)

irf_price_weak <- simulate_model(T, par_weak$alpha, par_weak$beta, par_weak$delta, par_weak$gamma,
                                 zp_path = zp1, x_path = x0) %>%
  mutate(case = "Weak feedback", experiment = "Price shock")

irf_price_strong <- simulate_model(T, par_strong$alpha, par_strong$beta, par_strong$delta, par_strong$gamma,
                                   zp_path = zp1, x_path = x0) %>%
  mutate(case = "Strong feedback", experiment = "Price shock")

# ---- IRF 2: Inflation response to permanent increase in x ----
zp0 <- rep(0, T)
x1  <- make_perm_x(T, shock_size = 1)

irf_x_weak <- simulate_model(T, par_weak$alpha, par_weak$beta, par_weak$delta, par_weak$gamma,
                             zp_path = zp0, x_path = x1) %>%
  mutate(case = "Weak feedback", experiment = "Permanent x increase")

irf_x_strong <- simulate_model(T, par_strong$alpha, par_strong$beta, par_strong$delta, par_strong$gamma,
                               zp_path = zp0, x_path = x1) %>%
  mutate(case = "Strong feedback", experiment = "Permanent x increase")

# ---- Combine and plot (like the slide) ----
irf_all <- bind_rows(irf_price_weak, irf_price_strong, irf_x_weak, irf_x_strong)

# Plot 1: price shock
p1 <- irf_all %>%
  filter(experiment == "Price shock") %>%
  ggplot(aes(x = quarter, y = inflation, colour = case)) +
  geom_line(linewidth = 1.1) +
  labs(
    title = "Figure 1. Responses of inflation to a price shock",
    x = "Quarter", y = "Percent", colour = NULL
  ) +
  theme_minimal()

# Plot 2: permanent x increase
p2 <- irf_all %>%
  filter(experiment == "Permanent x increase") %>%
  ggplot(aes(x = quarter, y = inflation, colour = case)) +
  geom_line(linewidth = 1.1) +
  labs(
    title = "Figure 2. Responses of inflation to a permanent increase in labor market tightness",
    x = "Quarter", y = "Percent", colour = NULL
  ) +
  theme_minimal()

print(p1)
print(p2)