# Take baseline model, for each variable compute pairwise significance of coef differences for each pair of years

dir.create('output/regressions/sig_ch')

# Get coefficients
coefs = read_csv('output/regressions/all_coefs.csv')

# Make 
ztest = function(v1,v2){
  (v1$coef - v2$coef) / sqrt(v1$se^2 + v2$se^2)
}


# Function to conduct z-tests for coefficient differences
compare_coefficients <- function(data, coef_col = "coefficient", se_col = "standard_error", 
                                 year_col = "year", alpha = 0.05) {
  
  # Input validation
  if (!coef_col %in% names(data) || !se_col %in% names(data) || !year_col %in% names(data)) {
    stop("Specified column names not found in data")
  }
  
  n <- nrow(data)
  if (n < 2) {
    stop("Need at least 2 rows to compare")
  }
  
  # Extract coefficients, standard errors, and years
  coefs <- data[[coef_col]]
  ses <- data[[se_col]]
  years <- data[[year_col]]
  
  # Initialize results storage
  results <- data.frame(
    year_1 = numeric(),
    year_2 = numeric(),
    coef_1 = numeric(),
    se_1 = numeric(),
    coef_2 = numeric(),
    se_2 = numeric(),
    difference = numeric(),
    se_difference = numeric(),
    z_statistic = numeric(),
    p_value = numeric(),
    significant = logical(),
    stringsAsFactors = FALSE
  )
  
  # Conduct pairwise comparisons
  row_counter <- 1
  for (i in 1:(n-1)) {
    for (j in (i+1):n) {
      
      # Calculate difference and its standard error
      diff <- coefs[j] - coefs[i]
      se_diff <- sqrt(ses[i]^2 + ses[j]^2)  # Assuming independence
      
      # Calculate z-statistic
      z_stat <- diff / se_diff
      
      # Calculate two-tailed p-value
      p_val <- 2 * (1 - pnorm(abs(z_stat)))
      
      # Determine significance
      is_sig <- p_val < alpha
      
      # Store results
      results[row_counter, ] <- list(
        year_1 = years[i],
        year_2 = years[j],
        coef_1 = coefs[i],
        se_1 = ses[i],
        coef_2 = coefs[j],
        se_2 = ses[j],
        difference = diff,
        se_difference = se_diff,
        z_statistic = z_stat,
        p_value = p_val,
        significant = is_sig
      )
      
      row_counter <- row_counter + 1
    }
  }
  
  return(results)
}


# Function to compare coefficients across multiple variables
compare_multiple_variables <- function(data, variable_col = "variable", 
                                       coef_col = "coefficient", 
                                       se_col = "standard_error", 
                                       year_col = "year", 
                                       alpha = 0.05) {
  
  # Input validation
  required_cols <- c(variable_col, coef_col, se_col, year_col)
  missing_cols <- required_cols[!required_cols %in% names(data)]
  if (length(missing_cols) > 0) {
    stop(paste("Missing columns:", paste(missing_cols, collapse = ", ")))
  }
  
  # Get unique variables
  variables <- unique(data[[variable_col]])
  
  # Initialize list to store results for each variable
  all_results <- list()
  
  # Loop through each variable
  for (var in variables) {
    cat("Processing variable:", var, "\n")
    
    # Subset data for current variable
    var_data <- data[data[[variable_col]] == var, ]
    
    # Check if we have enough observations for this variable
    if (nrow(var_data) < 2) {
      cat("  Warning: Variable", var, "has fewer than 2 observations. Skipping.\n")
      next
    }
    
    # Run comparison for this variable
    tryCatch({
      var_results <- compare_coefficients(var_data, 
                                          coef_col = coef_col, 
                                          se_col = se_col,
                                          year_col = year_col,
                                          alpha = alpha)
      
      # Add variable identifier to results
      var_results$variable <- var
      
      # Store results
      all_results[[var]] <- var_results
      
      cat("  Completed", nrow(var_results), "comparisons for", var, "\n")
      
    }, error = function(e) {
      cat("  Error processing variable", var, ":", e$message, "\n")
    })
  }
  
  # Combine all results
  if (length(all_results) > 0) {
    combined_results <- do.call(rbind, all_results)
    
    # Reorder columns to put variable first
    col_order <- c("variable", names(combined_results)[names(combined_results) != "variable"])
    combined_results <- combined_results[, col_order]
    
    # Reset row names
    rownames(combined_results) <- NULL
    
    return(combined_results)
  } else {
    cat("No valid comparisons were performed.\n")
    return(data.frame())
  }
}
  



# Do this for everything
coefs$incsamp_mod = paste(coefs$incsamp, coefs$mod,coefs$racedef,coefs$outvar)

for(v in unique(coefs$incsamp_mod)){
  
  chsig = coefs %>% filter(incsamp_mod == v)
  outtest = compare_multiple_variables(chsig, variable_col = 'var',coef_col = 'coef', se_col = 'se', year = 'year')
  
  
  write_csv(outtest, paste0('output/regressions/sig_ch/sig_test_',v,'.csv'))
  
  # Make figure
  heatmap_data = outtest #%>% filter(variable == 'pct_b_pop')
  heatmap_data = heatmap_data %>% mutate(
    sig_label = case_when(
      p_value < .05 & difference < 0 ~ 'Statistically significant decrease (p < .05)',
      p_value < .05 & difference > 0 ~ 'Statistically significant increase (p < .05)',
      .default = NA,
    )
  )
  
  fullvars = c("(Intercept)","logpop10",
               "pct_b_pop","pct_o_pop",  "pct_all_cit_imm",
               "union_combined","cemp_pct_manuf","cemp_pct_public",
               "dissimilarity_index","pct_all_col_ba",'pct_ba_lograt_bw','hrrv',
               'incar_comb_total','incar_comb_bwlrat',
               "regionMidwest","regionSouth","regionWest", 
               'incar_comb_bwdif','total_prison_pop_rate','imprate_vera_lrat_bw',
               'dissimilarity_index_ltdb','iso_black','iso_black_ltdb','exposure_b_w','exposure_b_w_ltdb','theil','theil_ltdb',
               'slave_pct:hrrvalid','cropfarms_pctblack:hrrvalid','hrrbaker:hrrvalid','hrr4:hrrvalid','numdis:hrrvalid','southman:hrrvalid','hrr_st:hrrvalid','hrr4_st:hrrvalid'
  )
  
  cleanvars = c('Intercept','MSA population (log10)', 
                '% Black','% neither Black nor white','% foreign-born',
                'Unionization rate','% Manufacturing','% Public administration',
                'BW dissimilarity index',"% with bachelor's degree",'Logged BA ratio','HRR factor',
                'State incarceration rate','Logged incarc. ratio',
                'Midwest','South','West',
                'BW difference in incarc.','State imprisionment rate','Logged imp. ratio',
                'BW dissimilarity index','Black isolation index','Black isolation index','BW exposure index','BW exposure index','BW Theil index','BW Theil index',
                '% Enslaved, 1860','% Black\nsharecroppers, 1930','HRR factor','HRR factor','Number of\ndisfranchisement devices','Support for\nSouthern Manifesto','HRR factor','HRR factor'
  )
  
  
  heatmap_data$clean_names = factor(heatmap_data$variable, levels = (fullvars), 
                                    labels = (cleanvars))
  
  
  p <- ggplot(heatmap_data, aes(y = factor(year_1), x = factor(year_2), fill = sig_label)) +
    geom_tile(color = "white", size = 0.5) +
    facet_wrap(.~clean_names, ncol = 4) + 
    geom_text(aes(label = z_statistic %>% round(2)), color = "black", size = 2, fontface = "bold") + #)
    # scale_fill_gradient2(low = "blue", mid = "white", high = "red", 
    scale_fill_discrete(type = c('red','blue'),na.value = NA, na.translate = F,#labels = c('','Statistically significant at p = 0.05'),
                        name = "Change from first year to second year") +
    labs(
      title = paste('Z-tests of coefficient changes over time'),
      # subtitle = v,
      x = "Second year",
      y = "First year",
      caption = "Positive values: Coefficient increased over time\nNegative values: Coefficient decreased over time"
    ) +
    theme_minimal() +
    guides(fill = guide_legend(ncol = 1)) + 
    theme(
      plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),
      plot.subtitle = element_text(hjust = 0.5, size = 10),
      plot.caption = element_text(hjust = 0.5, size = 8, color = "gray50"),
      axis.text.x = element_text(angle = 45, hjust = 1),
      panel.grid = element_blank(),legend.title.position = 'top',
      legend.position = "bottom"
    ) + scale_y_discrete(limits = rev) +
    coord_fixed()  # Make squares actually square
  
  print(p)
  ggsave(plot = p, filename = paste0('output/regressions/sig_ch/coef_ch_',v,'.pdf'), height = 12, width = 10)
  
  
}


# Now do all three in big one

incsss = c('all_faminc_sqrt Baseline expansive med_ldiff',
'all_faminc_sqrt Baseline expansive med_b',
'all_faminc_sqrt Baseline expansive med_w'
)

outfull = NULL

for(ins in incsss){
   
  outtest = read_csv(paste0('output/regressions/sig_ch/sig_test_',ins,'.csv'))
  outtest$incsamp_mod = ins
  
  outtest = outtest %>% mutate(
    sig_label = case_when(
      p_value < .05 & difference < 0 ~ 'Statistically significant decrease (p < .05)',
      p_value < .05 & difference > 0 ~ 'Statistically significant increase (p < .05)',
      .default = NA,
    )
  )
  
  outtest$clean_names = factor(outtest$variable, levels = (fullvars), 
                                    labels = (cleanvars))
  
  
  outfull = bind_rows(outfull, outtest)
}


inames = c('A. Median income ratio','B. Median income, Black families','C. Median income, white families')

# Create a temporary plot with legend to extract it
temp_plot <- ggplot(outfull[outfull$incsamp_mod == incsss[1], ], 
                    aes(y = factor(year_1), x = factor(year_2), fill = sig_label)) +
  geom_tile() +
  scale_fill_discrete(type = c('red','blue'), na.value = NA, na.translate = F,
                      name = "Change from first year to second year") +
  guides(fill = guide_legend(ncol = 1)) + 
  theme_minimal() +
  theme(
    legend.title.position = 'top',
    legend.position = "bottom"
  )

# Extract the legend
get_legend <- function(myggplot){
  tmp <- ggplot_gtable(ggplot_build(myggplot))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  return(legend)
}

shared_legend <- get_legend(temp_plot)


plot_list <- list()


# Create a plot for each incsamp_mode subset
for(i in 1:length(incsss)) {
  inmod = incsss[i]
  iname = inames[i]
  mode_data <- outfull %>% filter(incsamp_mod == inmod)
  
  p_temp <- ggplot(mode_data, aes(y = factor(year_1), x = factor(year_2), fill = sig_label)) +
    geom_tile(color = "white", size = 0.5) +
    facet_wrap(.~clean_names, ncol = 3) + 
    geom_text(aes(label = z_statistic %>% round(2)), color = "black", size = 2, fontface = "bold") +
    scale_fill_discrete(type = c('red','blue'), na.value = NA, na.translate = F,
                        name = "Change from first year to second year") +
    labs(
      title = paste( iname),
      x = "Second year",
      y = "First year",
      caption = "Positive values: Coefficient increased over time\nNegative values: Coefficient decreased over time"
    ) +
    theme_minimal() +
    guides(fill = guide_legend(ncol = 1)) + 
    theme(
      plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),
      plot.subtitle = element_text(hjust = 0.5, size = 10),
      plot.caption = element_text(hjust = 0.5, size = 8, color = "gray50"),
      axis.text.x = element_text(angle = 45, hjust = 1),
      panel.grid = element_blank(),
      legend.title.position = 'top',
      legend.position = "none"
    ) + 
    scale_y_discrete(limits = rev) +
    coord_fixed()
  
  plot_list[[i]] <- p_temp
}

# Combine all plots using patchwork or gridExtra
print(paste("Number of plots created:", length(plot_list)))
print(paste("Plot classes:", sapply(plot_list, class)))


# More robust approach - create grobs first
plot_grobs <- lapply(plot_list, ggplotGrob)

# Combine plots horizontally 
combined_plots <- do.call(arrangeGrob, c(plot_grobs, nrow = 1))

# Create the final layout with shared legend
final_plot <- arrangeGrob(
  combined_plots,
  shared_legend,
  nrow = 2,
  
  heights = unit(c(0.9, 0.1), "npc")   # More explicit height specification
)


ggsave(plot = final_plot, 
       filename = paste0('output/regressions/coef_ch_three_panels.pdf'), 
       width = 24, height =16)  # Increased height for three vertical panels



