## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment  = "#>",
  fig.width  = 7,
  fig.height = 4.5
)

## ----install, eval=FALSE------------------------------------------------------
# # Install dependencies
# install.packages(c("etwfe", "fixest", "ggplot2", "scales"))
# 
# # Install selectTWFE (once on CRAN)
# install.packages("selectTWFE")
# 
# # Or from GitHub
# devtools::install_github("your-username/selectTWFE")

## ----mpdta, message=FALSE, warning=FALSE--------------------------------------
library(selectTWFE)

data("mpdta", package = "selectTWFE")

result_mpdta <- select_twfe(
  fml    = lemp ~ lpop,
  tvar   = year,
  gvar   = first.treat,
  data   = mpdta,
  ivar   = countyreal,
  cgroup = "never",
  vcov   = ~countyreal
)

print(result_mpdta)

## ----mpdta-plot, message=FALSE, warning=FALSE---------------------------------
plot(result_mpdta)

## ----sim-homogeneous, message=FALSE, warning=FALSE----------------------------
set.seed(42)
n_units   <- 500
n_periods <- 6

# Assign cohorts: treated in period 3, 4, or 5; some never treated
cohort  <- sample(c(3, 4, 5, NA), n_units, replace = TRUE,
                  prob = c(.25, .25, .25, .25))
unit_fe <- rnorm(n_units, sd = 1)
time_fe <- rnorm(n_periods, sd = 0.5)

panel <- expand.grid(unit = 1:n_units, period = 1:n_periods)
panel$cohort  <- cohort[panel$unit]
panel$unit_fe <- unit_fe[panel$unit]
panel$time_fe <- time_fe[panel$period]
panel$treated <- !is.na(panel$cohort) & panel$period >= panel$cohort

# Homogeneous treatment effect = 1.0 for all cohorts and periods
panel$y    <- panel$unit_fe + panel$time_fe +
              1.0 * panel$treated + rnorm(nrow(panel), sd = 0.5)
panel$gvar <- ifelse(is.na(panel$cohort), Inf, panel$cohort)

result_homo <- select_twfe(
  fml    = y ~ 1,
  tvar   = period,
  gvar   = gvar,
  data   = panel,
  ivar   = unit,
  cgroup = "never",
  vcov   = ~unit
)

print(result_homo)

## ----sim-heterogeneous, message=FALSE, warning=FALSE--------------------------
set.seed(42)

# Substantial heterogeneous treatment effects: 0.5, 1.0, 2.0 by cohort
effect_map <- c("3" = 0.5, "4" = 1.0, "5" = 2.0)
effects <- ifelse(
  is.na(panel$cohort), 0,
  effect_map[as.character(panel$cohort)]
)
panel$y <- panel$unit_fe + panel$time_fe +
           effects * panel$treated + rnorm(nrow(panel), sd = 0.5)

result_hetero <- select_twfe(
  fml    = y ~ 1,
  tvar   = period,
  gvar   = gvar,
  data   = panel,
  ivar   = unit,
  cgroup = "never",
  vcov   = ~unit
)

print(result_hetero)
plot(result_hetero)

## ----sim-mild, message=FALSE, warning=FALSE-----------------------------------
set.seed(42)
n_units_small <- 30  # Very small sample

# Assign cohorts and build panel with small sample
cohort_small  <- sample(c(3, 4, 5, NA), n_units_small, replace = TRUE,
                        prob = c(.25, .25, .25, .25))
unit_fe_small <- rnorm(n_units_small, sd = 1)
time_fe_small <- rnorm(n_periods, sd = 0.5)

panel_small <- expand.grid(unit = 1:n_units_small, period = 1:n_periods)
panel_small$cohort  <- cohort_small[panel_small$unit]
panel_small$unit_fe <- unit_fe_small[panel_small$unit]
panel_small$time_fe <- time_fe_small[panel_small$period]
panel_small$treated <- !is.na(panel_small$cohort) & panel_small$period >= panel_small$cohort

# Very mild heterogeneous treatment effects
effect_map_veryMild <- c("3" = 0.90, "4" = 1.0, "5" = 1.10)
effects_veryMild <- ifelse(
  is.na(panel_small$cohort), 0,
  effect_map_veryMild[as.character(panel_small$cohort)]
)
panel_small$y <- panel_small$unit_fe + panel_small$time_fe +
                 effects_veryMild * panel_small$treated + rnorm(nrow(panel_small), sd = 0.5)
panel_small$gvar <- ifelse(is.na(panel_small$cohort), Inf, panel_small$cohort)

result_mild <- select_twfe(
  fml    = y ~ 1,
  tvar   = period,
  gvar   = gvar,
  data   = panel_small,
  ivar   = unit,
  cgroup = "never",
  vcov   = ~unit
)

print(result_mild)