library(tidyverse)
library(haven)
library(fixest)

mj <- data.frame(State = c("CO", "WA", "DC", "AK", "OR", "CA", "NV", "ME", "A", "VT", "MI", "IL"), Treatment = c(2012, 2012, 2014, 2014, 2014, 2016, 2016, 2016, 2016, 2018, 2018, 2019))

#this is the cleaned version of the UCR's arrest file
#via ICPSR/Jacob Kaplan
#https://www.openicpsr.org/openicpsr/project/102263/version/V15/view
#arrest <- read_dta("ucr_arrests_yearly_alcohol_or_property_crimes_age_1974_2020.dta")

#effect on possession and sales
state_cannabis_arrests <- arrest %>%
  filter(!arrest$state_abb %in% c("0", "CZ", "GU", "PR")) %>%
  group_by(year, state_abb) %>%
  summarise(sale_cannabis = sum(sale_cannabis_tot_arrests),
            poss_cannabis = sum(poss_cannabis_tot_arrests),
            population = sum(population)) %>%
  left_join(mj, by = c('state_abb' = 'State'))  %>%
  mutate(Treated = ifelse(is.na(Treatment), F, year > Treatment))

etable(feglm(data = state_cannabis_arrests, fml = sale_cannabis ~ Treated | year + state_abb),
       feglm(data = state_cannabis_arrests, fml = poss_cannabis ~ Treated | year + state_abb),
       feglm(data = state_cannabis_arrests, fml = log(sale_cannabis) ~ Treated | year + state_abb),
       feglm(data = state_cannabis_arrests, fml = log(poss_cannabis) ~ Treated | year + state_abb))

#effect on total arrests
total_arrest_cols <- arrest %>%
  select(ends_with('tot_arrests')) %>%
  #these appear to be redundant
  select(-gamble_total_tot_arrests, -total_drug_tot_arrests, -poss_drug_total_tot_arrests, -sale_drug_total_tot_arrests) %>%
  colnames()

total_arrests <- arrest %>%
  filter(!arrest$state_abb %in% c("0", "CZ", "GU", "PR")) %>%
  select(year, state_abb, population, matches(total_arrest_cols)) %>%
  group_by(year, state_abb) %>%
  summarise(total_arrests = sum(across(ends_with("tot_arrests"))),
            population = sum(population)) %>%
  left_join(mj, by = c('state_abb' = 'State'))  %>%
  mutate(Treated = ifelse(is.na(Treatment), F, year > Treatment))

#effect on black arrests

black_arrest_cols <- arrest %>%
  select(ends_with('tot_black')) %>%
  #these appear to be redundant
  select(-gamble_total_tot_black, -total_drug_tot_black, -poss_drug_total_tot_black, -sale_drug_total_tot_black) %>%
  colnames()

black_arrests <- arrest %>%
  filter(!arrest$state_abb %in% c("0", "CZ", "GU", "PR")) %>%
  select(year, state_abb, population, matches(black_arrest_cols)) %>%
  group_by(year, state_abb) %>%
  summarise(total_black_arrests = sum(across(ends_with("tot_black"))),
            population = sum(population)) %>%
  left_join(mj, by = c('state_abb' = 'State'))  %>%
  mutate(Treated = ifelse(is.na(Treatment), F, year > Treatment))

etable(feglm(data = total_arrests, fml = total_arrests ~ Treated | year + state_abb),
       feglm(data = black_arrests, fml = total_black_arrests ~ Treated | year + state_abb),
       feglm(data = total_arrests, fml = log(total_arrests) ~ Treated | year + state_abb),
       feglm(data = black_arrests, fml = log(total_black_arrests) ~ Treated | year + state_abb))