# Set seed for reproducibility
library(sandwich)
library(lmtest)

set.seed(20200403)

# Create the base data frame
n_states <- 40
n_firms_per_state <- 25
n_years <- 4

# Create expanded data frame for states and firms
states_rep <- rep(1:n_states, each = n_firms_per_state)
firms_data <- data.frame(
  state = states_rep,
  firms = runif(n_states * n_firms_per_state, 0, 5)
)

# Now expand for years
firms_data <- firms_data[rep(seq_len(nrow(firms_data)), each = n_years), ]
firms_data <- firms_data[order(firms_data$state), ]

# Create year variable
firms_data$year <- rep(1:4, times = nrow(firms_data)/4)
firms_data$n <- firms_data$year

# Replace years with actual dates
firms_data$year <- ifelse(firms_data$year == 1, 1990,
                          ifelse(firms_data$year == 2, 1991,
                                 ifelse(firms_data$year == 3, 1992, 1993)))

# Create unique ID for each firm
firms_data$id <- as.numeric(factor(paste(firms_data$state, firms_data$firms)))

# Treatment group assignment (upper half of IDs)
firms_data$group <- ifelse(firms_data$id >= 500, 1, 0)

# Create post indicators
firms_data$post <- ifelse(firms_data$year >= 1991, 1, 0)
firms_data$post_na <- ifelse(firms_data$year >= 1992, 1, 0)

# Generate error term
firms_data$e <- rnorm(nrow(firms_data), 0, 1)

# Generate potential outcomes
firms_data$y0 <- firms_data$firms + firms_data$n + firms_data$e

# Constant treatment effects
firms_data$y1_c <- firms_data$y0
firms_data$y1_c[firms_data$year >= 1991] <- 
  firms_data$y0[firms_data$year >= 1991] + 10

# Dynamic treatment effects
firms_data$y1_d <- firms_data$y0
firms_data$y1_d[firms_data$year == 1991] <- 
  firms_data$y0[firms_data$year == 1991] + 10
firms_data$y1_d[firms_data$year == 1992] <- 
  firms_data$y0[firms_data$year == 1992] + 20
firms_data$y1_d[firms_data$year == 1993] <- 
  firms_data$y0[firms_data$year == 1993] + 30

# Calculate treatment effects
firms_data$delta_c <- firms_data$y1_c - firms_data$y0
firms_data$delta_d <- firms_data$y1_d - firms_data$y0

# Create treatment indicator
firms_data$d <- ifelse(firms_data$year >= 1991 & firms_data$group == 1, 1, 0)

# Generate observed outcomes
firms_data$y_c <- firms_data$d * firms_data$y1_c + (1 - firms_data$d) * firms_data$y0
firms_data$y_d <- firms_data$d * firms_data$y1_d + (1 - firms_data$d) * firms_data$y0

# Calculate aggregate treatment effects
att_c <- mean(firms_data$delta_c[firms_data$year >= 1991 & firms_data$group == 1], 
              na.rm = TRUE)
att_d <- mean(firms_data$delta_d[firms_data$year >= 1991 & firms_data$group == 1], 
              na.rm = TRUE)

# Print summary of treatment effects
cat("Average Treatment Effects:\n")
cat("Constant ATT:", att_c, "\n")
cat("Dynamic ATT:", att_d, "\n\n")

# Fit regression models
# Correct specification - Constant Treatment Effects
model_c <- lm(y_c ~ factor(group) * factor(post), data = firms_data)
coeftest(model_c, vcov = vcovHC(model_c, type = "HC1"))

# Correct specification - Dynamic Treatment Effects
model_d <- lm(y_d ~ factor(group) * factor(post), data = firms_data)
coeftest(model_d, vcov = vcovHC(model_d, type = "HC1"))

# Incorrect specification - Constant Treatment Effects
model_c_na <- lm(y_c ~ factor(group) * factor(post_na), 
                 data = subset(firms_data, year >= 1991))
coeftest(model_c_na, vcov = vcovHC(model_c_na, type = "HC1"))

# Incorrect specification - Dynamic Treatment Effects
model_d_na <- lm(y_d ~ factor(group) * factor(post_na), 
                 data = subset(firms_data, year >= 1991))
coeftest(model_d_na, vcov = vcovHC(model_d_na, type = "HC1"))