noalign <- function(x) {
    x <- tinytable::theme_tinytable(x)
    fn <- function(table) {
        if (table@output != "typst") {
            return(table)
        }
        tab <- unlist(strsplit(table@table_string, "\\n"))
        idx <- grepl("^\\s*#align\\(center, \\[\\s*$|^\\s*\\]\\) // end align\\s*$", tab)
        table@table_string <- paste(tab[!idx], collapse = "\n")
        return(table)
    }
    x <- tinytable::style_tt(x, finalize = fn)
    return(x)
}
options(tinytable_tt_theme = noalign)

library("did")
library("broom")
library("fixest")
library("ggplot2")
library("data.table")

simulation6 = function() {
  dat = CJ(firm = 1:1000, year = 1980:2015)     [
    , time_fe := rnorm(1, sd = .5), by = "year"][
    , unit_fe := rnorm(1, sd = .5), by = "firm"][
    , state := sample(1:50, 1), by = "firm"    ]

  setkey(dat, state, firm, year)

  treatment_groups = data.table(
    state = c(1, 18, 35),
    cohort = c(1989, 1998, 2007),
    hat_gamma = c(.5, .3, .1))
  dat = treatment_groups[dat, roll = TRUE, on = "state"]

  dat                                                [
    , treat  := as.numeric(year >= cohort)          ][
    , gamma  := rnorm(.N, mean = hat_gamma, sd = .2)][
    , tau    := fifelse(treat == 1, gamma, 0)       ][
    , cumtau := cumsum(tau), by = "firm"            ][
    , error  := rnorm(.N, 0, .5)                    ][
    , y := unit_fe + time_fe + cumtau + error       ][
    , time_to_treat := year - cohort                ]

  return(dat)
}

dat = simulation6()

etwfe = feols(y ~ treat : factor(time_to_treat) : factor(cohort) | firm + year, data = dat)

# Clean the results
etwfe = as.data.table(tidy(etwfe))
etwfe = etwfe[
    , .(term = term, etwfe = estimate)][
    , group := as.numeric(gsub(".*cohort.", "", term))][
    , year := as.numeric(gsub(".*time_to_treat.(\\d+).*", "\\1", term)) + group][
    , .(group, year, etwfe)]

csa = att_gt(
    yname = "y",
    gname = "cohort",
    idname = "firm",
    tname = "year",
    control_group = "notyettreated",
    data = dat)

# Clean the results
csa = data.table(group = csa$group, year = csa$t, csa = csa$att)

# merge the TWFE and CSA results
results = merge(etwfe, csa, by = c("group", "year"))
colnames(results) = c("Cohort", "Year", "TWFE w/ interactions", "CSA (2021)")
results[, Cohort := factor(Cohort)]

dat_plot = melt(results, id.vars = c("Cohort", "Year"))
ggplot(dat_plot, aes(Year, value, color = variable, linetype = Cohort)) +
  geom_line(size = 1.4) +
  theme_minimal() +
  labs(x = "Year", y = "ATT", color = "Estimator", linetype = "Cohort")

ggplot(results, aes(`TWFE w/ interactions`, `CSA (2021)`, color = Cohort)) +
  geom_point(size = 2) +
  geom_abline(intercept = 0, slope = 1) +
  labs(title = "On the 45 degree line, estimates of the group-time ATT\nare identical under the two strategies.") +
  theme_minimal()

library(fixest)
library(broom)
library(tidyverse)
library(ggokabeito)
theme_set(theme_minimal())
scale_colour_discrete <- scale_colour_okabe_ito
scale_fill_discrete <- scale_fill_okabe_ito

# set seed
set.seed(20200403)
# Fixed Effects ------------------------------------------------
# unit fixed effects
unit <- tibble(
unit = 1:1000, 
unit_fe = rnorm(1000, 0, 1),
# generate state
state = sample(1:40, 1000, replace = TRUE),
# generate treatment effect
mu = rnorm(1000, 0.3, 0.2))

# year fixed effects 
year <- tibble(
year = 1980:2010,
year_fe = rnorm(31, 0, 1))

# Trend Break -------------------------------------------------------------
# Put the states into treatment groups
treat_taus <- tibble(
  # sample the states randomly
  state = sample(1:40, 40, replace = FALSE),
  # place the randomly sampled states into five treatment groups G_g
  cohort_year = sort(rep(c(1982, 1991, 1998, 2004), 10)))

# make main dataset
# full interaction of unit X year 
dat <- expand_grid(unit = 1:1000, year = 1980:2010) %>% 
  left_join(., unit, by = "unit") %>% 
  left_join(., year, by = "year") %>% 
  left_join(., treat_taus, by = "state") %>% 
  # make error term and get treatment indicators and treatment effects
  mutate(error = rnorm(31000, 0, 0.5),
         treat = ifelse(year >= cohort_year, 1, 0),
         tau = ifelse(treat == 1, mu, 0)) %>% 
  # calculate cumulative treatment effects
  group_by(unit) %>% 
  mutate(tau_cum = cumsum(tau)) %>% 
  ungroup() %>% 
  # calculate the dep variable
  mutate(dep_var = unit_fe + year_fe + tau_cum + error,
         w = as.numeric(year >= cohort_year))

mod <- feols(dep_var ~
             w : i(cohort_year, i.year) +
             i(year) |
             unit,
             data = dat)

cohorts <- c(as.character(sort(unique(dat$cohort_year))), NA)

res <- mod %>%
    # extract
    tidy(conf.int = TRUE) %>%
    # ignore year dummies
    filter(grepl("^w", term)) %>%
    # cleanup
    separate(term, sep = "::", into = c(NA, "cohort", "year")) %>%
    mutate(cohort = gsub(":.*", "", cohort),
           year = as.numeric(year),
           time_to_treatment = year - as.numeric(cohort)) %>%
    # ATTs are available for a cohort only until the next cohort gets treated
    mutate(cohort_next = as.numeric(cohorts[match(cohort, cohorts) + 1])) %>%
    filter(year < cohort_next)

ggplot(res, aes(x = time_to_treatment,
                y = estimate,
                ymin = conf.low,
                ymax = conf.high,
                color = cohort)) +
    # Baker simulation:
    # tau is a random variable equal with mean 0.3
    # cumsum(tau) = ATT
    geom_abline(intercept = .3, slope = .3) +
    geom_pointrange() +
    labs(x = "Time to treatment",
         y = "ATT",
         color = "Cohort",
         title = "Group-Time ATTs estimated using TWFE & Interactions.",
         subtitle = "The black line shows the truth.")