#clear the environment 
rm(list=ls()) 
#setwd(path)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
# install.packages(c("ggplot2","dplyr","tidyr"))
# Optional (nicer multi-panel layout): install.packages("patchwork")
library(ggplot2)
library(dplyr)
library(tidyr)

has_patchwork <- requireNamespace("patchwork", quietly = TRUE)
if (has_patchwork) library(patchwork)

# -----------------------------
# 1) Parameters
# -----------------------------
beta0 <- 1.0
beta1 <- 0.9
ty    <- 0.30
R     <- 0.02

T <- 100

shock_size <- 0.5   # shock to g via epsilon_t
shock_time <- 1     # shock happens in period t = 1

B_init_lag <- 0     # initial lagged debt: B_{-1}

# -------------------------------------------------------
# 1b) Steady-state initialization for t_{-1}
#     (consistent with lagged-tax spending rule and B_{-1})
#     y* = (beta0 - R B*) / ((1-beta1)(1-ty))
#     t* = ty * y*
# -------------------------------------------------------
B_ss <- B_init_lag
y_ss <- (beta0 - R * B_ss) / ((1 - beta1) * (1 - ty))
t_init_lag <- ty * y_ss   # set t_{-1} = t*

# -------------------------------------------------------
# 2) Within-period equilibrium given (B_{t-1}, t_{t-1}, eps_t)
#    c_t = beta0 + beta1 (y_t - t_t)
#    t_t = ty y_t
#    g_t = t_{t-1} - R B_{t-1} + eps_t   <-- reacts to lagged taxes
#    y_t = c_t + g_t
# -------------------------------------------------------
solve_period <- function(B_lag, t_lag, eps, beta0, beta1, ty, R) {
  # y = beta0 + beta1(1-ty)y + (t_lag - R B_lag + eps)
  denom <- 1 - beta1 * (1 - ty)
  if (abs(denom) < 1e-10) stop("Denominator too close to zero; check parameters.")
  
  y <- (beta0 + t_lag - R * B_lag + eps) / denom
  t_tax <- ty * y
  g <- t_lag - R * B_lag + eps
  c <- y - g
  
  list(c = c, t_tax = t_tax, g = g, y = y)
}

# ----------------------------------------
# 3) Simulation
#    B_t = (1+R)B_{t-1} + (g_t - t_t)
#    and next period uses t_{t-1} = t_t
# ----------------------------------------
simulate_model <- function(T, eps_path, beta0, beta1, ty, R, B_init_lag, t_init_lag) {
  out <- tibble(
    t = 0:(T-1),
    eps = eps_path,
    c = NA_real_, t_tax = NA_real_, g = NA_real_, y = NA_real_, B = NA_real_
  )
  
  B_lag <- B_init_lag
  t_lag <- t_init_lag
  
  for (i in 1:T) {
    sol <- solve_period(B_lag, t_lag, out$eps[i], beta0, beta1, ty, R)
    
    out$c[i]     <- sol$c
    out$t_tax[i] <- sol$t_tax
    out$g[i]     <- sol$g
    out$y[i]     <- sol$y
    
    out$B[i] <- (1 + R) * B_lag + (out$g[i] - out$t_tax[i])
    
    # update states for next period
    B_lag <- out$B[i]
    t_lag <- out$t_tax[i]
  }
  
  out
}

# -----------------------------
# 4) Scenarios
# -----------------------------
eps_base  <- rep(0, T)

eps_shock <- rep(0, T)
eps_shock[shock_time + 1] <- shock_size  # t=0 is index 1

base  <- simulate_model(T, eps_base,  beta0, beta1, ty, R, B_init_lag, t_init_lag) %>%
  mutate(scenario = "Baseline")
shock <- simulate_model(T, eps_shock, beta0, beta1, ty, R, B_init_lag, t_init_lag) %>%
  mutate(scenario = "Shock")

paths <- bind_rows(base, shock)

# -----------------------------
# (Optional) quick sanity check:
# baseline should keep g - t_tax = 0 and B constant at B_init_lag
# -----------------------------
print(base %>% transmute(t, g, t_tax, gap = g - t_tax, B))

# -----------------------------
# 5) Plots of levels
# -----------------------------
ggplot(paths, aes(x = t, y = g, linetype = scenario)) +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Government Spending (g)",
       x = "Time", y = "Level") +
  theme_minimal() +
  theme(legend.position = "top")

ggplot(paths, aes(x = t, y = y, linetype = scenario)) +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Output (y)",
       x = "Time", y = "Level") +
  theme_minimal() +
  theme(legend.position = "top")

ggplot(paths, aes(x = t, y = B, linetype = scenario)) +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Government Debt (B)",
       x = "Time", y = "Level") +
  theme_minimal() +
  theme(legend.position = "top")

# -----------------------------
# 6) Plots of IRFs
# -----------------------------
irf <- shock %>%
  select(t, g, y, B) %>%
  left_join(base %>% select(t, g, y, B),
            by = "t", suffix = c("_shock","_base")) %>%
  mutate(
    irf_g = g_shock - g_base,
    irf_y = y_shock - y_base,
    irf_B = B_shock - B_base
  )

ggplot(irf, aes(x = t, y = irf_g)) +
  geom_hline(yintercept = 0, linetype = "dashed") +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Impulse Response of g",
       x = "Time", y = "Shock - Baseline") +
  theme_minimal()

ggplot(irf, aes(x = t, y = irf_y)) +
  geom_hline(yintercept = 0, linetype = "dashed") +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Impulse Response of y",
       x = "Time", y = "Shock - Baseline") +
  theme_minimal()

ggplot(irf, aes(x = t, y = irf_B)) +
  geom_hline(yintercept = 0, linetype = "dashed") +
  geom_vline(xintercept = shock_time, linetype = "dotted") +
  geom_line(linewidth = 1.2) +
  labs(title = "Impulse Response of B",
       x = "Time", y = "Shock - Baseline") +
  theme_minimal()