# Simulations to decompose contribution of changing covariates vs changing coefficients

dir.create('output/simulations')

# Bring in outcomes and control variabes
outv = read_csv('intermediate_data/regression_data/msa_reg_outcomes.csv') %>% select(-st,-region)
controlv = read_csv('intermediate_data/regression_data/msa_reg_predictors.csv')

# Merge
fulldta = full_join(outv, controlv, by = c('year','racedef','msa'))
stopifnot(nrow(fulldta) == nrow(outv))
stopifnot(!is.na(fulldta$pop_all_full))

# Make south variable
fulldta$south = fulldta$region == 'South'
fulldta$region = factor(fulldta$region, levels = c('Northeast','Midwest','South','West'))

incsamps = setdiff(unique(outv$incsamp),'emp_lf_faminc_sqrt') #c('all_faminc_sqrt') #unique(outv$incsamp) ,'all_hhinc_sqrt','all_faminc_norm','all_faminc_tot','edu_hs_faminc_sqrt','edu_ba_faminc_sqrt','cit_usborn_faminc_sqrt','emp_emp_inctot' ,'gq_hh_faminc_sqrt','edu_over25_faminc_sqrt'
yrs = unique(fulldta$year)
rdefs = unique(fulldta$racedef)


# Models - run baseline
models = c()
modnames = c()

i = 1

modnames[i] = 'Baseline'
models[i] = 'outvar ~ logpop10 +
  pct_b_pop + pct_o_pop + pct_all_cit_imm + pct_all_col_ba + pct_ba_lograt_bw +
  cemp_pct_manuf + cemp_pct_public + incar_comb_total + incar_comb_bwlrat +
  union_combined + hrrv + dissimilarity_index'


modnum = i
mod = models[modnum] %>% formula()
modname = modnames[modnum]


# Run baseline model and sample
incs = 'all_faminc_sqrt'
rdef = 'expansive'
outv = 'med_ldiff'

# Get sample data - consistent 1969 sample

rdta = filter(fulldta, incsamp == incs, racedef == rdef,pop_b > 0)

# Limit to just places observed in the 1969 regression and with incarceration observed
segmsas = filter(rdta, year == 1969,!is.na(dissimilarity_index)) %>% select(msa) %>% unlist() %>% unique()
incarmsas = filter(rdta, !is.na(incar_comb_bwlrat), year == 2017) %>% select(msa) %>% unlist()
msas_consist = intersect(segmsas,incarmsas)


rdta = filter(rdta, msa %in% msas_consist)

rdta = rdta %>% rename(outvar = !!sym(outv))

dta69 = rdta %>% filter(year == 1969)

# Get consistent sample
dta17 = rdta %>% filter(year == 2017)


# Run for 1969
lm1 = lm_robust(mod,data = dta69, weights = pop_b )

summary(lm1) # Matches output for consistent MSA sample


# Now try standardizing first
modvars = c('outvar',setdiff(names(lm1$coefficients),c('(Intercept)','regionMidwest','regionSouth','regionWest','southTRUE')))

# Have to scale manually because it's population-weighted
sdta69 = mutate(dta69, across(all_of(modvars),~ (.x - wtd.mean(.x,pop_b)) / sqrt(wtd.var(.x,pop_b)) %>% as.vector()))
sdta17 = mutate(dta17, across(all_of(modvars),~ (.x - wtd.mean(.x,pop_b)) / sqrt(wtd.var(.x,pop_b)) %>% as.vector()))



# Run again to see how it looks
lm2 = lm_robust(mod, data = sdta69, weights = pop_b )

summary(lm2)  # Matches standardized coefs from main output


# Get 2017 coefficients
lm3 = lm_robust(mod,  #+ pct_all_gq + log(lynchings + 1)
         data = dta17, weights = pop_b )

summary(lm3) # Matches output

lm4 = lm_robust(mod,  #+ pct_all_gq + log(lynchings + 1)
         data = sdta17, weights = pop_b )

summary(lm4) # Matches standardized output



## Get predicted values and residuals
sdta17$pred_69coefs17dist = predict(lm2,newdata = sdta17)
sdta17$pred_17coefs17dist = predict(lm4,newdata = sdta17)
sdta17$resid_17coefs17dist = sdta17$outvar - sdta17$pred_17coefs17dist
sdta69$pred_17coefs69dist = predict(lm4,newdata = sdta69)
sdta69$pred_69coefs69dist = predict(lm2,newdata = sdta69)
sdta69$resid_69coefs69dist = sdta69$outvar - sdta69$pred_69coefs69dist



sdta = left_join(sdta17,sdta69 %>% select(msa,pred_17coefs69dist,pred_69coefs69dist,resid_69coefs69dist,outvar,all_of(modvars),south), by = 'msa',suffix = c('17','69'))

sdta = mutate(sdta, 
              rpred_69coefs17dist69resid = pred_69coefs17dist + resid_69coefs69dist, # Old resid, old coefs, new dist
              rpred_69coefs17dist17resid = pred_69coefs17dist + resid_17coefs17dist, # New resid, old coefs, new dist
              rpred_69coefs69dist17resid = pred_69coefs69dist + resid_17coefs17dist, # New resid, old coefs, old dist
              rpred_17coefs69dist69resid = pred_17coefs69dist + resid_69coefs69dist, # Old resid, new coefs, old dist
              rpred_17coefs69dist17resid = pred_17coefs69dist + resid_17coefs17dist, # New resid, new coefs, old dist
              rpred_17coefs17dist69resid = pred_17coefs17dist + resid_69coefs69dist, # Old resid, new coefs, new dist
              ch_resid = resid_17coefs17dist - resid_69coefs69dist, # Direct change in resid
              ch_obs = outvar17 - outvar69, # Overall change
              
              ch_resid_first = rpred_69coefs69dist17resid - outvar69, # Change from resid first - should match direct change in resid
              ch_resid_second = rpred_69coefs17dist17resid - rpred_69coefs17dist69resid, # Change from resid second - should match chagne from resid first
              ch_resid_last = outvar17 - rpred_17coefs17dist69resid, # Change from resid last - should match other resid changes
              
              ch_dist_first = rpred_69coefs17dist69resid - outvar69, # Change from dist first - should match change from dist after resid
              ch_dist_secondrfirst = rpred_69coefs17dist17resid - rpred_69coefs69dist17resid, # Change from dist after resid - should match change from dist before resid
              
              ch_coefs_seconddfirst = rpred_17coefs17dist69resid - rpred_69coefs17dist69resid, # Change from all coefs after dist before resid - should match next one and sum of individual coefs
              ch_coefs_last = outvar17 - rpred_69coefs17dist17resid, # Change from all coefs after dist and resid - should match previous
              
              #Now do coefs before dist
              ch_coefs_first = rpred_17coefs69dist69resid - outvar69 ,# coefficients first
              ch_coefs_secondrfirst = rpred_17coefs69dist17resid - rpred_69coefs69dist17resid ,# Coefficients after residuals before dist - should match coefs first
              
              ch_dist_secondcfirst = rpred_17coefs17dist69resid - rpred_17coefs69dist69resid, # Dist after coefficients before resid
              ch_dist_last = outvar17 - rpred_17coefs69dist17resid # Dist last
)

# Now collapse into South and not
southd = sdta %>% filter(!is.na(resid_69coefs69dist)) %>% group_by(south17) %>% summarise(
  across(c(all_of(c('outvar17','outvar69','pred_69coefs17dist','pred_17coefs69dist','rpred_69coefs17dist69resid','rpred_69coefs17dist17resid',
                  'rpred_17coefs69dist69resid','rpred_17coefs69dist17resid','rpred_69coefs69dist17resid','rpred_17coefs17dist69resid',paste0(modvars,'69'),paste0(modvars,'17'))),starts_with('ch_'),starts_with('coef')),~ wtd.mean(.x, pop_b))

)

southd = southd %>% rename(south = south17)


# Now get long version of just changes
slong = southd %>% pivot_longer(-south)
slong = slong %>% pivot_wider(id_cols = name,names_from = south,values_from = value,names_prefix = 'south')
slong$regdif = slong$southTRUE - slong$southFALSE

write_csv(slong,'output/simulations/slong_std.csv')



# Now do just percent of change but do exploded stacked bar
pdta = filter(slong, (name %in% c('ch_dist_first','ch_resid_second','ch_coefs_last')) | (str_detect(name,'coefch_dist17')))


pdta$name = factor(pdta$name, levels = c('outvar69','ch_dist_first','ch_coefs_last','ch_resid_second'),
                   labels = c('Observed 1969 value','Change in pop. chars.','Change in coefficients','Change in MSA-level residuals'
                   ))

pdta = pdta %>% arrange(name) %>% filter(!is.na(name))


pdta = mutate(pdta, 
              sf_pct = southFALSE / sum(southFALSE),
              st_pct = southTRUE / sum(southTRUE),
              dif_pct = regdif / sum(regdif)
)


pdta$endpoint = cumsum(pdta$dif_pct)
pdta$startpoint = pdta$endpoint - pdta$dif_pct


pdf('output/simulations/chart_simulations_explode_17dist_std.pdf', height = 5,width = 5)
gplt = ggplot(pdta %>% filter(name != "Observed 1969 value"), aes(x = name,xend = name, y = startpoint, yend = endpoint, group = 1)) + 
  theme(axis.text.x = element_text(hjust = 1,angle = 50)) +
  geom_hline(lty = 2, yintercept = 0, alpha = .6) + 
  geom_hline(lty = 3, yintercept = 1, alpha = .6) + 
  geom_segment(linewidth = 1,lineend = 'square',linejoin = 'mitre',arrow = NULL) + 
  scale_x_discrete(name = '') + 
  geom_segment(aes(y = endpoint -.005, yend = endpoint + .005), linewidth = 6) +
  geom_segment(aes(y = startpoint -.005, yend = startpoint + .005), linewidth = 6) +
  scale_y_continuous(name = 'Fraction of total change in difference\n between the South and the rest of country', labels = percent)
print(gplt)
dev.off()
write_csv(pdta, 'output/simulations/dta_explode_17_std.csv')


# Now do 1969 dist

pdta = filter(slong, (name %in% c('ch_dist_secondcfirst','ch_resid_second','ch_coefs_first'))) 

pdta$name = factor(pdta$name, levels = c('outvar69','ch_dist_secondcfirst','ch_coefs_first','ch_resid_second'),
                   labels = c('Observed 1969 value','Change in pop. chars.','Change in coefficients','Change in MSA-level residuals'))
                   
pdta = pdta %>% arrange(name) %>% filter(!is.na(name))


pdta = mutate(pdta, 
              sf_pct = southFALSE / sum(southFALSE),
              st_pct = southTRUE / sum(southTRUE),
              dif_pct = regdif / sum(regdif)
)


pdta$endpoint = cumsum(pdta$dif_pct)
pdta$startpoint = pdta$endpoint - pdta$dif_pct


pdf('output/simulations/chart_simulations_explode_69dist_std.pdf', height = 5,width = 5)
gplt = ggplot(pdta %>% filter(name != "Observed 1969 value"), aes(x = name,xend = name, y = startpoint, yend = endpoint, group = 1)) + 
  theme(axis.text.x = element_text(hjust = 1,angle = 50)) +
  geom_hline(lty = 2, yintercept = 0, alpha = .6) + 
  geom_hline(lty = 3, yintercept = 1, alpha = .6) + 
  geom_segment(linewidth = 1,lineend = 'square',linejoin = 'mitre',arrow = NULL) +  
  scale_x_discrete(name = '') + 
  geom_segment(aes(y = endpoint -.005, yend = endpoint + .005), linewidth = 6) +
  geom_segment(aes(y = startpoint -.005, yend = startpoint + .005), linewidth = 6) +
  scale_y_continuous(name = 'Fraction of total change in difference\n between the South and the rest of country', labels = percent)
print(gplt)
dev.off()

write_csv(pdta, 'output/simulations/dta_explode_69_std.csv')

# Get 
dta69 %>% group_by(south) %>% summarise(out69 = wtd.mean(outvar, pop_b))
dta17 %>% group_by(south) %>% summarise(out17 = wtd.mean(outvar, pop_b))

-0.542 - (-0.733) # Narrowing of 0.191 in the South
-0.646 - (-0.517) # Widening of -0.129 in the South
-0.129 - 0.191 # Total change is 0.32 log points

exp(-0.542) # 0.5815839 in 2017 in South
exp(-0.733) # 0.4804654 in 1969 in South
exp(-0.646) # 0.5241381 in 2017 in North
exp(-0.517) # 0.5963068 in 1969 in North



