############################
# MAIN ANALYSES
#
#
# Part of replication files for:
# Dill, Janina, Howlett, Marnie, & Müller-Crepon, Carl (2025).
# Do Ukrainians Still Prefer Self-Defense Against Russia At Any Cost? 
# Journal of Peace Research, forthcoming.
#
#############################


# AMCE and MargMeans by Experiment ####
for(out in outvars){
  for(s in samples){
    print(s)
    coef.df <- do.call(rbind, lapply(0:2, function(e){
      cbind(experiment = e, 
            rbind(est_acme(df = data.long[data.long$experiment == e & data.long[,s], ], 
                           outvar = out, attr.vec = seq_along(attr.labs),
                           attr.labs = attr.labs, refvals = rep(1, length(attr.labs)), coadj = FALSE),
                  est_margmeans(df = data.long[data.long$experiment == e & data.long[,s], ], 
                                outvar = out, attr.vec = seq_along(attr.labs),
                                attr.labs = attr.labs)))
    }))
    coef.df$experiment_lab <- c("0 (2022-07)", "I (2024-12/2025-01)", "II (2024-12/2025-01)")[coef.df$experiment + 1]
    
    # Prepare plot
    coef.df <- prep_plot(coef.df)
    
    # Means
    means <- unique(coef.df[, c("treat_type","est"), drop = F])
    means$mean = ifelse(means$est == "AMCE", 0, .5)
    coef.df$ypos <- as.numeric(coef.df$treat_lab) + .25 - coef.df$experiment/4
    
    # Plot
    g <- ggplot(coef.df, 
                aes(x = coef, y = ypos, col = factor(experiment))) +
      geom_vline(data = means, aes(xintercept = mean, x = NULL, y = NULL), lty = 2, col = "grey") +
      geom_point() +
      geom_segment(aes(yend = ypos, 
                       x = coef - 1.96*se, 
                       xend = coef + 1.96*se)) +
      facet_grid(treat_type ~ est,
                 scales  = "free", space = "free_y",switch = "y") +
      theme_minimal() +
      labs(y = NULL, x = "Estimate", col = "Experiment") +
      scale_color_discrete(labels = unique(coef.df$experiment_lab)) +
      scale_y_continuous(breaks = as.numeric(coef.df$treat_lab) ,
                         labels = as.character(coef.df$treat_lab),
                         position = "right") +
      theme(panel.spacing.x = unit(1, "lines"),
            panel.spacing.y = unit(1, "lines"),
            strip.text.y.right = element_text(angle = 0, hjust = 0, size = 5),
            panel.grid.major.y = element_blank(),
            panel.grid.minor.y = element_line(linewidth = .5),
            legend.position = "top")
    
    
    ## Save
    png(file.path(fig.path, paste0("all_experiments_main",s, 
                                   ifelse(out == "choice", "", "_score"),".png")),
        width = 6.5, height = 6.5, units = "in", res = 400)
    print(g)
    dev.off()
    
  }
}

# SIGNIFICANT DIFFERENCE?
data.long$exp0vsexp1_22 <- ifelse(data.long$experiment == 0, "Exp 0", 
                                  ifelse(data.long$experiment == 1 & data.long$in_2022 == 1, "Exp 1", NA))
data.long$exp0vsexp2_22 <- ifelse(data.long$experiment == 0, "Exp 0", 
                                  ifelse(data.long$experiment == 2 & data.long$in_2022 == 1, "Exp 2", NA))
data.long$exp1vsexp2_22 <- ifelse(data.long$experiment == 1, "Exp 1", 
                                  ifelse(data.long$experiment == 2 & data.long$in_2022 == 1, "Exp 2", NA))
data.long$exp0vsexp1_24 <- ifelse(data.long$experiment == 0, "Exp 0", 
                                  ifelse(data.long$experiment == 1 & data.long$in_2024 == 1, "Exp 1", NA))
data.long$exp0vsexp2_24 <- ifelse(data.long$experiment == 0, "Exp 0", 
                                  ifelse(data.long$experiment == 2 & data.long$in_2024 == 1, "Exp 2", NA))
data.long$exp1vsexp2_24 <- ifelse(data.long$experiment == 1, "Exp 1", 
                                  ifelse(data.long$experiment == 2 & data.long$in_2024 == 1, "Exp 2", NA))

all.df <- heteff_acme(df = data.long,
                      outvars = outvars,
                      hetvars = list(c(`Experiment 0 vs 1, Sample A` = "exp0vsexp1_22", 
                                       `Experiment 0 vs 2, Sample A` = "exp0vsexp2_22", 
                                       `Experiment 1 vs 2, Sample A` = "exp1vsexp2_22",
                                       `Experiment 0 vs 1, Sample B` = "exp0vsexp1_24", 
                                       `Experiment 0 vs 2, Sample B` = "exp0vsexp2_24", 
                                       `Experiment 1 vs 2, Sample B` = "exp1vsexp2_24")),
                      attr.vec = 1:5,
                      attr.labs = attr.labs, prefix = paste0("expcomp"), 
                      summary = FALSE, plot = FALSE)
unique(all.df[,c("hetvar","wald_pval", "wald_stat")])


## PValues for single differences

### Experiment 1 vs 2 // Nuclear Risk
summary(feols(choice ~ factor(attr4_level)*exp1vsexp2_24, 
              data = data.long[!is.na(data.long$exp1vsexp2_24),],
              cluster = "id"))

### Experiment 0 vs 2 // Nuclear Risk
summary(feols(choice ~ factor(attr4_level)*exp0vsexp2_24, 
              data = data.long[!is.na(data.long$exp0vsexp2_24),],
              cluster = "id"))

### Experiment 0 vs 2 // Military fatalities
summary(feols(choice ~ factor(attr3_level)*exp0vsexp1_24, 
              data = data.long[!is.na(data.long$exp0vsexp1_24),],
              cluster = "id"))
summary(feols(choice ~ factor(attr3_level)*exp0vsexp2_24, 
              data = data.long[!is.na(data.long$exp0vsexp2_24),],
              cluster = "id"))




# AMCEs CONDITIONAL ON ATTRIBUTES 1 AND 5 ###########
# Figure 5
test.ls <- list()
for(o in outvars){
  for(s in samples["Sample B"]){
    print(s)
    # Estimation
    for(e in c(as.list(1:2), list(c(1:2)))){
      
      print(o)
      intvars <- c("attr1_level", "attr5_level")
      coef.df <- do.call(rbind, lapply(intvars, function(i){
        
        # Locals /variables, labels, etc)
        if((i == "attr1_level")){
          attr.vars <- paste0("attr",rep(2:5, 2), "_d", rep(2:3, each = 4))
          all.treat.vars <- paste0("attr",rep(2:5, 3), "_d", rep(1:3, each = 4))
          treat_lab = as.vector(do.call(rbind, attr.labs[2:5]))
          treat_type = rep(names(attr.labs)[2:5],3)
          int_lab <- as.vector(do.call(rbind, attr.labs[1]))
        } else {
          attr.vars <- paste0("attr",rep(1:4, 2), "_d", rep(2:3, each = 4))
          all.treat.vars <- paste0("attr",rep(1:4, 3), "_d", rep(1:3, each = 4))
          treat_lab = as.vector(do.call(rbind, attr.labs[1:4]))
          treat_type = rep(names(attr.labs)[1:4],3)
          int_lab <- as.vector(do.call(rbind, attr.labs[5]))
        }
        int.dummies <- sapply(c("d1", "d2", "d3"), function(s){gsub("level", s, i)})
        
        # Formula
        f <- as.formula(paste(o, " ~ ",  paste(int.dummies[-1], 
                                               collapse = "+"), 
                              " + ", 
                              paste(paste0("I(",
                                           rep(attr.vars, each = 3), "*",
                                           rep(paste0("(",i, "==", 1:3, ")"), 2), ")"),
                                    collapse = " + ")))
        # Estimate
        m <- fixest::feols(f,
                           data = data.long[data.long$experiment %in% e & data.long[,s], ],
                           cluster = clustvar)
        
        # F-Test of significance of interactions
        
        ## Reestimate as interction with constitutive terms
        restr.vars <- paste0("I(",
                             rep(attr.vars[grepl("attr2|attr3|attr4",attr.vars)], each = 2), "*",
                             rep(paste0("(",i, "==", 2:3, ")"), 2), ")")
        mf <-  as.formula(paste(o, " ~ ",  "factor(", i, ") + ", 
                                paste(c(attr.vars[grepl("attr2|attr3|attr4",attr.vars)], 
                                        restr.vars),
                                      collapse = " + ")))
        
        ## Estimate
        mf <- fixest::feols(mf,
                            data = data.long[data.long$experiment %in% e & data.long[,s], ],
                            cluster = clustvar)
        
        ## Test
        test <- wald(mf, 
                     keep = paste0(i, " =="), 
                     cluster = clustvar,
                     print = T)
        test.ls <<- c(test.ls, list(test))
        
        # Extract all results
        res.df <- rbind(data.frame(outcome = names(outvars)[outvars == o],
                                   intvar = i,
                                   intvar_val = rep(1:3),
                                   level = rep(1:3),
                                   treat = int.dummies,
                                   treat_lab = int_lab,
                                   treat_type = names(attr.labs)[!names(attr.labs) %in% treat_type],
                                   coef = c(0, m$coefficients[2:3]),
                                   se = c(NA, m$se[2:3]),
                                   p = c(NA, summary(m)$coeftable[2:3,4]),
                                   wald_stat = test$stat, wald_pval = test$p),
                        data.frame(outcome = names(outvars)[outvars == o],
                                   intvar = i,
                                   intvar_val = rep(1:3, 4*3),
                                   level = rep(1:3, each = 4*3),
                                   treat = rep(all.treat.vars, each = 3),
                                   treat_lab = rep(treat_lab, each = 3),
                                   treat_type = rep(treat_type, each = 3),
                                   coef = c(rep(0, 4*3), m$coefficients[-c(1:3)]),
                                   se = c(rep(NA, 4*3), m$se[-c(1:3)]),
                                   p = c(rep(NA, 4*3), summary(m)$coeftable[-c(1:3),4]),
                                   wald_stat = test$stat, wald_pval = test$p)
        )
        
        # Return
        return(res.df)
        
      }))
      
      # Prepare plot
      coef.df$outcome <- factor(coef.df$outcome, 
                                levels = names(outvars),
                                ordered = T)
      coef.df$treat_type <- factor(coef.df$treat_type, 
                                   levels = names(attr.labs),
                                   ordered = T)
      coef.df$treat_lab <- factor(coef.df$treat_lab, 
                                  levels = rev(unique(unlist(attr.labs))),
                                  ordered = T)
      coef.df$int_type <- ifelse(coef.df$intvar == "attr1_level",
                                 "Territorial Integrity", "Political Autonomy")
      coef.df$int_type <- factor(coef.df$int_type, levels = unique(coef.df$int_type),
                                 ordered = T)
      
      coef.df$yval <- as.numeric(coef.df$treat_lab) +
        -1*(coef.df$intvar_val - 2)/4
      coef.df$yval <- ifelse(gsub("\n", " ", as.character(coef.df$treat_type)) == coef.df$int_type,
                             as.numeric(coef.df$treat_lab),
                             coef.df$yval)
      
      # Plot
      g <- ggplot(coef.df[coef.df$treat_type %in% names(attr.labs)[-c(1,5)], ], 
                  aes(x = coef, y = yval, 
                      col = factor(intvar_val), 
                      pch = factor(intvar_val))) +
        geom_vline(xintercept = 0, lty = 2, col = "grey") +
        geom_point() +
        geom_segment(aes(yend = yval, 
                         x = coef - 1.96*se, 
                         xend = coef + 1.96*se)) +
        facet_grid(treat_type ~ int_type,
                   scales  = "free",
                   switch = "y") +
        scale_y_continuous(breaks = as.numeric(coef.df$treat_lab) ,
                           labels = as.character(coef.df$treat_lab) ,
                           position = "right") +
        theme_minimal() +
        labs(y = NULL,
             x = paste0("Conditional Average Marginal Component Effect on ", names(outvars)[outvars == o]),
             col = "Level of territorial integrity (left) or political automy (right):", pch = "Level of territorial integrity (left) or political automy (right):" ) +
        theme(panel.spacing = unit(.75, "lines"),
              panel.grid.major.y = element_blank(),
              # strip.text.y.left = element_text(angle = 0, hjust = 0),
              panel.grid.minor.y = element_line(size = .5),
              legend.position = "top",
              legend.box.margin=margin(-10,-10,-10,-10),
              plot.caption = element_text(hjust = 0, colour = "grey20")) +
        scale_color_manual(values = custom.col)
      
      # Save plot
      png(file.path(fig.path, paste0("exp", paste(e, collapse = ""), s, "_moder_amce_",o,".png")),
          width = 7, height = 4, units = "in", res = 400)
      print(g)
      dev.off()

    }
  }
}

# "UNAMBIGUOUS" PAIRS #######

# Make Variable
data.long$ambiguity <- ifelse(data.long$attr1_amb == 1 & data.long$attr5_amb == 1,
                              "Invariant attributes 1 and 5", "Variant attribute 1 or 5" )
data.long$ambiguity <- factor(data.long$ambiguity, 
                              levels = c("Variant attribute 1 or 5", "Invariant attributes 1 and 5"),
                              ordered = T)

# Values
hvals <- sort(na.omit(unique(data.long[,"ambiguity"])))
h <- "ambiguity"
test.ls <- list()

# Estimate
#   This makes only sense w.r.t. the choice outcome, since the
#   scores are unrestricted in a pair -- both can be good or bad. 
for(s in samples["Sample B"]){
  for(e in c(as.list(1:2), list(1:2))){
    coef.df <- do.call(rbind, lapply(outvars[2], function(o){
      
      # Locals
      attr.vars <- paste0("attr",rep(2:4, 2), "_d", rep(2:3, each = 3))
      all.treat.vars <- paste0("attr",rep(2:4, 3), "_d", rep(1:3, each = 3))
      treat_lab = as.vector(do.call(rbind, attr.labs[2:4]))
      treat_type = rep(names(attr.labs)[2:4],3)
      
      # Formula
      f <- as.formula(paste(o, " ~ ",  "factor(", h, ") + ", 
                            paste(paste0("I(",
                                         rep(attr.vars, each = length(hvals)), "*",
                                         rep(paste0("(",h, "=='", hvals, "')"), 2), ")"),
                                  collapse = " + ")))
      # Estimate
      m <- fixest::feols(f,
                         data = data.long[data.long$experiment %in% e & data.long[,s], ],
                         cluster = clustvar)
      
      # F-Test of significance of interactions
      
      ## Re-estimate as interaction
      restr.vars <- paste0("I(",
                           rep(attr.vars, each = length(hvals) - 1), "*",
                           rep(paste0("(",h, "=='", hvals, "')"), 2), ")")
      rf <-  as.formula(paste(o, " ~ ",  "factor(", h, ") + ", 
                              paste(c(attr.vars, 
                                      restr.vars),
                                    collapse = " + ")))
      mf <- fixest::feols(rf,
                          data = data.long[data.long$experiment %in% e & data.long[,s], ],
                          cluster = clustvar)
      
      ## F-Test
      test <- wald(mf, 
                   keep = paste0(h, " =="), 
                   cluster = clustvar, print = F)
      test.ls <<- c(test.ls, list(test))
      # Extract results
      drop <- 1:length(hvals)
      res.df <- data.frame(outcome = names(outvars)[outvars == o],
                           hetvar = h,
                           hetvar_val = as.numeric(rep(hvals, 3*3)),
                           hetvar_lab = rep(hvals, 3*3),
                           level = rep(1:3, each = 3*length(hvals)),
                           treat = rep(all.treat.vars, each = length(hvals)),
                           treat_lab = rep(treat_lab, each = length(hvals)),
                           treat_type = rep(treat_type, each = length(hvals)),
                           coef = c(rep(0, 3*length(hvals)), m$coefficients[-drop]),
                           se = c(rep(NA, 3*length(hvals)), m$se[-drop]),
                           p = c(rep(NA, 3*length(hvals)), summary(m)$coeftable[-drop,4]),
                           wald_stat = test$stat,
                           wald_pval = test$p)
      
      # Return
      return(res.df)
    }))
    
    # Prepare plot
    coef.df$outcome <- factor(coef.df$outcome, 
                              levels = names(outvars),
                              ordered = T)
    coef.df$treat_type <- factor(coef.df$treat_type, 
                                 levels = names(attr.labs),
                                 ordered = T)
    coef.df$treat_lab <- factor(coef.df$treat_lab, 
                                levels = rev(unique(unlist(attr.labs))),
                                ordered = T)
    coef.df$het_type <- coef.df$hetvar
    
    coef.df$yval <- as.numeric(coef.df$treat_lab) +
      -1*((coef.df$hetvar_val - mean(coef.df$hetvar_val)) /
            max(coef.df$hetvar_val))/1.5
    
    # Plot
    g <- ggplot(coef.df, 
                aes(x = coef, y = yval, 
                    col = factor(hetvar_lab), 
                    pch = factor(hetvar_lab))) +
      geom_vline(xintercept = 0, lty = 2, col = "grey") +
      geom_point() +
      geom_segment(aes(yend = yval, 
                       x = coef - 1.96*se, 
                       xend = coef + 1.96*se)) +
      facet_grid(treat_type ~ outcome,
                 scales  = "free",
                 switch = "y") +
      scale_y_continuous(breaks = as.numeric(coef.df$treat_lab) ,
                         labels = as.character(coef.df$treat_lab),
                         minor_breaks = seq(.5, max(as.numeric(coef.df$treat_lab)) + .5, by = 1) ,
                         position = "right") +
      xlim(c(-.325,.05)) +
      theme_minimal() +
      labs(y = NULL,
           x = paste0("Conditional Average Marginal Component Effect"),
           col = "Strategy pairs with:", pch = "Strategy pairs with:"
      ) +
      theme(panel.spacing = unit(1, "lines"),
            legend.position = "top",
            panel.grid.major.y = element_blank(),
            panel.grid.minor.y = element_line(size = .5)) +
      scale_color_manual(values = custom.col)
    
    # Save
    png(file.path(fig.path, paste0("exp",paste(e, collapse = ""), s, "_heter_amce_",h,".png")),
        width = 6.5, height = 4.11, units = "in", res = 400)
    print(g)
    dev.off()

  }
}

