#############################
# FUNCTIONS
#
#
# 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.
#
# All ranking functions are generalized here:
# https://github.com/carl-mc/cjRank
# 
#############################


#########################
# Estimation functions
#########################

est_acme <- function(df, outvar, attr.vec, attr.labs, refvals = NULL, weight.var = NULL, coadj = TRUE){
  # Setup
  attr.names.vec <- unlist(lapply(attr.labs, names))
  attr.lab.vec <- unlist(attr.labs)
  names(attr.lab.vec) <- attr.names.vec
  
  ## Reference points
  if(is.null(refvals)){
    refvals <- rep(1, length(attr.vec))
    ref.pts <- paste0("attr", attr.vec, "_d", refvals)
  } else {
    stopifnot(length(refvals) == length(attr.vec))
    ref.pts <- paste0("attr", attr.vec, "_d", refvals)
  }
  
  ### AMCE & MargMeans Variables
  mm.vars <- sort(colnames(df))
  amce.vars <- mm.vars <- 
    unlist(sapply(attr.vec, function(a){mm.vars[grepl(paste0("attr",a), mm.vars) & 
                                                  grepl("_d", mm.vars, fixed = T)]}))
  amce.vars <- amce.vars[!amce.vars %in% ref.pts]
  
  # Data
  df <- df[!is.na(df[,outvar]), , drop = F]
  if(nrow(df) == 0){
    warning("No data")
    return(NULL)
  }
  
  # Weights
  if(!is.null(weight.var)){
    weights <- df[,weight.var]
  } else {
    weights <- rep(1, nrow(df))
  }
  
  # Estimate
  if(coadj){
    # Estimation
    do.call(rbind, lapply(attr.vec, function(a){
      # print(a)
      # Which AMCE?
      thse.vars <- amce.vars[grepl(paste0("attr",a, "_d"), amce.vars)]
      
      # Anything to estimate?
      if(length(thse.vars) == 0){
        return(NULL)
      } else {
        # Formula
        f <- as.formula(paste(outvar, " ~ ", 
                              paste(thse.vars,
                                    collapse = " + ")))
        # Data
        df <- df[df[,paste0("attr", a, "_amb")] %in%
                   c(0, ifelse(outvar == "choice", 0, 1)) &
                   !is.na(df[,outvar]), , drop = F]
        if(nrow(df) == 0){
          return(NULL)
        }
        if(!is.null(weight.var)){
          weights <- df[,weight.var]
        } else {
          weights <- rep(1, nrow(df))
        }
        
        # Estimate
        m <- fixest::feols(f,
                           data = df,
                           cluster = clustvar,
                           weights = weights)
        
        # Extract results
        df <- rbind(data.frame(outcome = names(outvars)[outvars == outvar],
                               outvar = outvar,
                               attr =  a,
                               level = refvals[attr.vec == a],
                               treat = ref.pts[attr.vec == a],
                               coef = 0,
                               se = NA,
                               p = NA, est = "AMCE"),
                    data.frame(outcome = names(outvars)[outvars == outvar],
                               outvar = outvar,
                               attr = a,
                               level = as.numeric(substr(thse.vars, 
                                                         nchar(thse.vars), 
                                                         nchar(thse.vars))),
                               treat = thse.vars,
                               coef = m$coefficients[-1],
                               se = m$se[-1],
                               p = summary(m)$coeftable[-1,4], est = "AMCE"))
      }
      
      
      # Finalize
      df$treat_lab = attr.lab.vec[df$treat]
      df$treat_type = names(attr.labs)[a]
      df <- df[order(df$level),]
      df
    }))
  } else {
    # Formula
    f <- as.formula(paste(outvar, " ~ ", 
                          paste(amce.vars,
                                collapse = " + ")))
    # Estimate
    m <- fixest::feols(f,
                       data = df,
                       cluster = clustvar,
                       weights = weights)
    
    # Extract results
    df <- rbind(
      data.frame(outcome = names(outvars)[outvars == outvar],
                 outvar = outvar,
                 attr = as.numeric(gsub("\\_.*|attr", "", amce.vars)),
                 level = as.numeric(substr(amce.vars, nchar(amce.vars), nchar(amce.vars))),
                 treat = amce.vars,
                 coef = m$coefficients[-1],
                 se = m$se[-1],
                 p = summary(m)$coeftable[-1,4], est = "AMCE"),
      data.frame(outcome = names(outvars)[outvars == outvar],
                 outvar = outvar,
                 attr =  as.numeric(gsub("\\_.*|attr", "", mm.vars[!mm.vars %in% amce.vars])),
                 level = refvals,
                 treat = ref.pts,
                 coef = 0,
                 se = NA,
                 p = NA, est = "AMCE")
    )
    
    df$treat_lab = attr.lab.vec[df$treat]
    df$treat_type = names(attr.labs)[df$attr]
    df <- df[order(df$attr, df$level),]
    df
  }
}


est_margmeans <- function(df, outvar, attr.vec, attr.labs, weight.var = NULL, coadj = TRUE){
  # Setup
  attr.names.vec <- unlist(lapply(attr.labs, names))
  attr.lab.vec <- unlist(attr.labs)
  names(attr.lab.vec) <- attr.names.vec
  
  
  ### AMCE & MargMeans Variables
  mm.vars <- sort(colnames(df))
  amce.vars <- mm.vars <- 
    unlist(sapply(attr.vec, function(a){mm.vars[grepl(paste0("attr",a), mm.vars) & 
                                                  grepl("_d", mm.vars, fixed = T)]}))
  amce.vars <- amce.vars[!grepl("d1", amce.vars)]
  
  
  # Estimation
  do.call(rbind, lapply(attr.vec, function(a){
    # Formula
    f <- as.formula(paste(outvar, " ~ 0 + ", 
                          paste(mm.vars[grepl(paste0("attr",a, "_d"), mm.vars)],
                                collapse = " + ")))
    # Data
    df <- df[df[,paste0("attr", a, "_amb")] %in%
               c(0, ifelse(outvar == "choice" & coadj, 0, 1)) &
               !is.na(df[,outvar]), , drop = F]
    if(nrow(df) == 0){
      return(NULL)
    }
    if(!is.null(weight.var)){
      weights <- df[,weight.var]
    } else {
      weights <- rep(1, nrow(df))
    }
    
    # Estimate
    m <- fixest::feols(f,
                       data = df,
                       cluster = clustvar,
                       weights = weights)
    
    # Extract results
    df <- data.frame(outcome = names(outvars)[outvars == outvar],
                     outvar = outvar,
                     attr = a,
                     level = seq_along(mm.vars[grepl(paste0("attr",a, "_d"), mm.vars)]),
                     treat = mm.vars[grepl(paste0("attr",a, "_d"), mm.vars)],
                     coef = m$coefficients,
                     se = m$se,
                     p = summary(m)$coeftable[,4], est = "Marginal Means")
    df$treat_lab = attr.lab.vec[df$treat]
    df$treat_type = names(attr.labs)[a]
    df <- df[order(df$level),]
    df
  }))
}

prep_plot <- function(coef.df){
  coef.df$outcome <- factor(coef.df$outcome, 
                            levels = names(outvars),
                            ordered = T)
  
  fac.vars <- c("treat_type", "int_type","treat_lab","int_lab","sample")
  fac.vars <- fac.vars[fac.vars %in% colnames(coef.df)]
  for(v in fac.vars){
    if(v %in% c("treat_lab")){
      coef.df[,v] <- factor(coef.df[,v], 
                            levels = rev(unique(coef.df[,v])),
                            ordered = T)
    } else {
      coef.df[,v] <- factor(coef.df[,v], 
                            levels = unique(coef.df[,v]),
                            ordered = T)
    }
    
  }
  
  coef.df
}


# HETEROGENEOUS AMCE ESTIMATES ##############
heteff_acme <- function(df,  outvars, hetvars, attr.vec, attr.labs, prefix, summary = TRUE, plot = TRUE){
  # Setup
  attr.names.vec <- unlist(lapply(attr.labs, names))
  attr.lab.vec <- unlist(attr.labs)
  names(attr.lab.vec) <- attr.names.vec
  het.names.vec <- unlist(lapply(hetvars, names))
  
  # AMCE & MargMeans Variables
  mm.vars <- sort(colnames(df))
  amce.vars <- mm.vars <- 
    unlist(sapply(attr.vec, function(a){mm.vars[grepl(paste0("attr",a), mm.vars) & 
                                                  grepl("_d", mm.vars, fixed = T)]}))
  amce.vars <- amce.vars[!grepl("d1", amce.vars)]
  
  # Omnibus test results empty init
  omn.test.res <- NULL
  all.plot.df <- NULL
  # Estimation
  for(h in unlist(hetvars)){
    print(h)
    # Values
    hvals <- sort(na.omit(unique(df[,h])))
    
    # Estimate
    coef.df <- do.call(rbind, lapply(outvars, function(o){
      
      # Prepare estimation and saving of results
      attr.vars <- amce.vars
      all.treat.vars <- mm.vars
      # treat_lab = as.vector(do.call(rbind, attr.labs[1:5]))
      # treat_type = rep(names(attr.labs)[1:5],3)
      
      # Data
      df <- df[!is.na(df[,o]) & !is.na(df[,h]), , drop = F]
      if(nrow(df) == 0){
        warning("No data")
        return(NULL)
      }
      
      # 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 = df,
                         cluster = clustvar)
      
      # F-Test of significance of interactions
      
      ## Reformulate estimation function to normal interaction, i.e. w/ constitutive term
      restr.vars <- paste0("I(",
                           rep(attr.vars, each = length(hvals)-1), "*",
                           rep(paste0("(",h, "=='", hvals[-1], "')"), length(attr.vars)), ")")
      mf <-  as.formula(paste(o, " ~ ",  "factor(", h, ") + ", 
                              paste(c(attr.vars, 
                                      restr.vars),
                                    collapse = " + ")))
      
      ## Estimate
      mf <- fixest::feols(mf,
                          data = df,
                          cluster = clustvar)
      
      ## F-Test
      test <- wald(mf, 
                   keep = paste0(h, " =="), 
                   cluster = clustvar, print = F)
      
      ## Save for later
      omn.test.res <<- rbind(omn.test.res ,
                             data.frame(outcome =  names(outvars)[outvars == o],
                                        hetvar = h,
                                        wald_fstat = test$stat, 
                                        wald_pval = test$p))
      
      
      # Extract results
      res.df <- rbind(
        data.frame(outcome = names(outvars)[outvars == o],
                   outvar = o,
                   hetvar = h,
                   hetvar_val = as.numeric(rep(factor(hvals), length(amce.vars))),
                   hetname = het.names.vec[unlist(hetvars) == h],
                   hetvar_lab = rep(factor(hvals), length(amce.vars)),
                   attr = rep(as.numeric(gsub("\\_.*|attr", "", amce.vars)), each = length(hvals)),
                   level = rep(as.numeric(substr(amce.vars, nchar(amce.vars), nchar(amce.vars))), each = length(hvals)),
                   treat = rep(amce.vars, each = length(hvals)),
                   coef = m$coefficients[-seq_along(hvals)],
                   se = m$se[-seq_along(hvals)],
                   p = summary(m)$coeftable[-seq_along(hvals),4],
                   wald_stat = test$stat, 
                   wald_pval = test$p),
        data.frame(outcome = names(outvars)[outvars == o],
                   outvar = o,
                   hetvar = h,
                   hetname = het.names.vec[unlist(hetvars) == h],
                   hetvar_val = as.numeric(rep(factor(hvals), length(mm.vars[!mm.vars %in% amce.vars]))),
                   hetvar_lab = rep(factor(hvals), length(mm.vars[!mm.vars %in% amce.vars])),
                   attr =  rep(as.numeric(gsub("\\_.*|attr", "", mm.vars[!mm.vars %in% amce.vars])), each = length(hvals)),
                   level = rep(1, each = length(hvals)),
                   treat = rep(mm.vars[!mm.vars %in% amce.vars], each = length(hvals)),
                   coef = 0,
                   se = NA,
                   p = NA,
                   wald_stat = test$stat, 
                   wald_pval = test$p)
      )
      
      res.df$treat_lab = attr.lab.vec[res.df$treat]
      res.df$treat_type = names(attr.labs)[res.df$attr]
      res.df <- res.df[order(res.df$attr, res.df$level),]
      res.df
      
      ## Return
      return(res.df)
    }))
    
    # Combine all results for this variable and prepare plot
    coef.df <- prep_plot(coef.df)
    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)) / 2
    
    # print(coef.df)
    
    # Plot
    if(plot){
      het.name = het.names.vec[unlist(hetvars) == h]
      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(size = .85) +
        # geom_point(aes(y = yval + .15), col = "transparent") +
        # geom_point(aes(y = yval - .15), col = "transparent") +
        geom_segment(aes(yend = yval, 
                         x = coef - 1.96*se, 
                         xend = coef + 1.96*se)) +
        facet_grid(treat_type ~ outcome,
                   scales  = "free", space = "free_y",
                   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("Average Marginal Component Effect "),
             col = paste0(het.name,":"), 
             pch = paste0(het.name,":"),
             caption = paste0("Moderator variable: ", het.name, 
                              "\nRespondents: ",
                              paste(paste0(hvals, " (",sapply(hvals, function(x){sum(df[!duplicated(df$id), h] == x, na.rm = T)}), ")"),
                                    collapse = ", "),
                              "\nWald-test of nullity of sub-group differences in AMCEs:\n",
                              paste(sapply(unique(as.character(coef.df$outcome)), function(t){
                                paste0(t, " (F: ", 
                                       signif(coef.df$wald_stat[coef.df$outcome == t][1], 2), "; p: ",
                                       signif(coef.df$wald_pval[coef.df$outcome == t][1], 2), ")")
                              }), collapse = ", "))) +
        theme(panel.spacing = unit(.75, "lines"),
              strip.text.y.left = element_text(angle = 0, hjust = 0),
              legend.position = "top",
              panel.grid.major.y = element_blank(),
              panel.grid.minor.y = element_line(size = .5),
              legend.box.margin=margin(-10,-10,-10,-10),
              plot.caption = element_text(hjust = 0, colour = "grey20")) +
        scale_color_manual(values = custom.col)
      
      # Save
      png(file.path(fig.path, paste0(prefix, "_hetamce_",h,".png")),
          width = 7, height = 6, units = "in", res = 400)
      print(g)
      dev.off()
    } else {
      all.plot.df <- rbind(all.plot.df, coef.df)
    }
    
  }
  
  
  # OVERVIEW TABLE #################
  if(summary){
    # Adjust p-alue / Bonferroni
    omn.test.res$p.adjust <- stats::p.adjust(omn.test.res$wald_pval,
                                             method = "bonferroni")
    
    # Outcome factor
    omn.test.res$outcome <- factor(omn.test.res$outcome,
                                   levels = unique(omn.test.res$outcome),
                                   ordered = T)
    
    # Prepare table 
    
    ## Long -> wide by outcome
    tab <- do.call(cbind, split(omn.test.res[,c("wald_fstat", "wald_pval", "p.adjust")], 
                                omn.test.res$outcome))
    tab <- cbind(var = het.names.vec,
                 tab)
    
    ## Feed to stargazer
    tab.tex <- stargazer(tab,
                         title = "Omnibus Wald-Test Result for Joint Nullity of Heterogenous Effects by Moderator",
                         summary = FALSE,
                         column.sep.width = "0pt",
                         label = paste0("tab:", prefix, "heter_test"),
                         rownames = FALSE,
                         digits = 2, align = T,
                         notes.append = F, 
                         notes = paste0("Note: Adjusted p-values based on a Bonferroni adjustment for ", length(omn.test.res$wald_pval)," hypotheses."))
    
    ## Fix stargazer output
    tab.tex[11] <- paste0(paste(paste0(" \\multicolumn{1}{c}{",
                                       c("Moderator", rep(c("F-Stat", "$p$", "Adj. $p$"), 2)), "}"), 
                                collapse = "&"), "\\\\")
    tab.tex[9] <- paste(tab.tex[9], 
                        "& \\multicolumn{3}{c}{Score (0-1)} & \\multicolumn{3}{c}{Forced choice (0/1)} \\\\ \\\\[-1.8ex]")
    tab.tex[c(13 + c(0, cumsum(sapply(hetvars, length)[-length(hetvars)])))] <- paste0("\\\\[-1.8ex] \\multicolumn{1}{l}{\\textbf{",
                                                                                       names(hetvars),
                                                                                       "}} \\\\",
                                                                                       tab.tex[c(13 + c(0, cumsum(sapply(hetvars, length)[-length(hetvars)])))])
    tab.tex[13:length(tab.tex)] <- gsub("multicolumn{1}{c}", "multicolumn{1}{l}", 
                                        tab.tex[13:length(tab.tex)], fixed = T)
    
    # Write table to disc
    writeLines(tab.tex, 
               file.path(tab.path, paste0(prefix, "_het_test.tex")))
  }
  
  if(!plot){
    return(all.plot.df)
  }
  
}



# HETEROGENEOUS MARGINAL MEANS ESTIMATES ##############
heteff_mm <- function(df,  outvars, hetvars, attr.vec, attr.labs, prefix, add_amce = TRUE, plot = TRUE){
  # Setup
  attr.names.vec <- unlist(lapply(attr.labs, names))
  attr.lab.vec <- unlist(attr.labs)
  names(attr.lab.vec) <- attr.names.vec
  het.names.vec <- unlist(lapply(hetvars, names))
  
  # AMCE & MargMeans Variables
  mm.vars <- sort(colnames(df))
  amce.vars <- mm.vars <- 
    unlist(sapply(attr.vec, function(a){mm.vars[grepl(paste0("attr",a), mm.vars) & 
                                                  grepl("_d", mm.vars, fixed = T)]}))
  amce.vars <- amce.vars[!grepl("d1", amce.vars)]
  
  # Omnibus test results empty init
  all.plot.df <- NULL
  # Estimation
  for(h in unlist(hetvars)){
    print(h)
    # Values
    hvals <- sort(na.omit(unique(df[,h])))
    
    # Estimate
    coef.df <- do.call(rbind, lapply(outvars, function(o){
      
      # Prepare estimation and saving of results
      attr.vars <- amce.vars
      all.treat.vars <- mm.vars
      # treat_lab = as.vector(do.call(rbind, attr.labs[1:5]))
      # treat_type = rep(names(attr.labs)[1:5],3)
      
      # Data
      df <- df[!is.na(df[,o]) & !is.na(df[,h]), , drop = F]
      if(nrow(df) == 0){
        warning("No data")
        return(NULL)
      }
      
      # Estimate
      res.df <- do.call(rbind, lapply(hvals, function(hv){
        coef.df <- est_margmeans(df[df[,h] == hv, ], o, attr.vec, attr.labs, coadj = TRUE)
        
        if(add_amce){
          coef.df <- rbind(coef.df,
                           est_acme(df = df[df[,h] == hv, ], 
                                    outvar = o, attr.vec = attr.vec,
                                    attr.labs = attr.labs, 
                                    refvals = rep(1, length(attr.labs)), 
                                    coadj = FALSE))
        }
        coef.df$hetvar = h
        coef.df$hetvar_val = as.numeric(factor(hvals))[hvals == hv]
        coef.df$hetname = het.names.vec[unlist(hetvars) == h]
        coef.df$hetvar_lab = factor(hvals)[hvals == hv]
        coef.df
      }))
      
      res.df
      
      ## Return
      return(res.df)
    }))
    
    # Combine all results for this variable and prepare plot
    coef.df <- prep_plot(coef.df)
    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)) / 2
    
    # print(coef.df)
    # Means
    means <- unique(coef.df[, c("treat_type","est"), drop = F])
    means$mean = ifelse(means$est == "AMCE", 0, .5)
    
    if(!add_amce){
      means <- means[means$est != "AMCE", , drop = F]
    }
    
    # Plot
    if(plot){
      het.name = het.names.vec[unlist(hetvars) == h]
      g <- ggplot(coef.df, 
                  aes(x = coef, y = yval, 
                      col = factor(hetvar_lab), 
                      pch = factor(hetvar_lab))) +
        geom_vline(data = means, aes(xintercept = mean, x = NULL, 
                                     y = NULL), lty = 2, col = "grey") +
        geom_point(size = .85) +
        # geom_point(aes(y = yval + .15), col = "transparent") +
        # geom_point(aes(y = yval - .15), col = "transparent") +
        geom_segment(aes(yend = yval, 
                         x = coef - 1.96*se, 
                         xend = coef + 1.96*se)) +
        facet_grid(treat_type ~ est,
                   scales  = "free", space = "free_y",
                   switch = "y") +
        scale_y_continuous(breaks = as.numeric(coef.df$treat_lab) ,
                           labels = as.character(coef.df$treat_lab),
                           position = "right", expand = expansion(add = .35)) +
        theme_minimal() +
        labs(y = NULL,
             x = paste0("Estimate"),
             col = paste0(het.name,":"), 
             pch = paste0(het.name,":"),
             caption = paste0("Moderator variable: ", het.name, 
                              "\nRespondents: ",
                              paste(paste0(hvals, " (",sapply(hvals, function(x){sum(df[!duplicated(df$id), h] == x, na.rm = T)}), ")"),
                                    collapse = ", "),
                              # "\nWald-test of nullity of sub-group differences in AMCEs:\n",
                              # paste(sapply(unique(as.character(coef.df$outcome)), function(t){
                              #   paste0(t, " (F: ", 
                              #          signif(coef.df$wald_stat[coef.df$outcome == t][1], 2), "; p: ",
                              #          signif(coef.df$wald_pval[coef.df$outcome == t][1], 2), ")")
                              # }), collapse = ", "),
                              ""
             )) +
        theme(panel.spacing = unit(.75, "lines"),
              strip.text.y.left = element_text(angle = 0, hjust = 0),
              legend.position = "top",
              panel.grid.major.y = element_blank(),
              panel.grid.minor.y = element_line(size = .5),
              legend.box.margin=margin(-10,-10,-10,-10),
              plot.caption = element_text(hjust = 0, colour = "grey20")) +
        scale_color_manual(values = custom.col)
      
      # Save
      png(file.path(fig.path, paste0(prefix, "_",h,".png")),
          width = 6.5, height = 6.5, units = "in", res = 400)
      print(g)
      dev.off()
    } else {
      all.plot.df <- rbind(all.plot.df, coef.df)
    }
    
  }
  
  
  if(!plot){
    return(all.plot.df)
  }
  
}



# HETEROGENEOUS AMCE ESTIMATES ##############
heteff_acme_fsl <- function(df,  outvars, hetvars, attr.vec, attr.labs, fsl, refgrp, prefix, summary = TRUE, plot = TRUE){
  # Setup
  attr.names.vec <- unlist(lapply(attr.labs, names))
  attr.lab.vec <- unlist(attr.labs)
  names(attr.lab.vec) <- attr.names.vec
  het.names.vec <- unlist(lapply(hetvars, names))
  
  # AMCE & MargMeans Variables
  mm.vars <- sort(colnames(df))
  amce.vars <- mm.vars <- 
    unlist(sapply(attr.vec, function(a){mm.vars[grepl(paste0("attr",a), mm.vars) & 
                                                  grepl("_d", mm.vars, fixed = T)]}))
  amce.vars <- amce.vars[!grepl("d1", amce.vars)]
  
  # Omnibus test results empty init
  omn.test.res <- NULL
  all.plot.df <- NULL
  # Estimation
  for(h in unlist(hetvars)){
    print(h)
    # Values
    all.hvals <- sort(na.omit(unique(df[,h])))
    hvals <- all.hvals[!all.hvals %in% refgrp]
    
    # Estimate
    coef.df <- do.call(rbind, lapply(outvars, function(o){
      
      # Prepare estimation and saving of results
      attr.vars <- amce.vars
      all.treat.vars <- mm.vars
      # treat_lab = as.vector(do.call(rbind, attr.labs[1:5]))
      # treat_type = rep(names(attr.labs)[1:5],3)
      
      # Data
      df <- df[!is.na(df[,o]) & !is.na(df[,h]), , drop = F]
      if(nrow(df) == 0){
        warning("No data")
        return(NULL)
      }
      
      # Formula
      f <- as.formula(paste(o, " ~ ",  "factor(", h, ") + ", 
                            paste(paste0("I(",
                                         rep(attr.vars, each = length(hvals)), "*",
                                         rep(paste0("(",h, "=='", hvals, "')"), 2), ")"),
                                  collapse = " + "),
                            ifelse(!is.null(fsl),
                                   paste0(" | ", fsl),
                                   "")))
      # Estimate
      m <- fixest::feols(f,
                         data = df,
                         cluster = clustvar)
      
      # F-Test of significance of interactions
      
      ## Reformulate estimation function to normal interaction, i.e. w/ constitutive term
      restr.vars <- paste0("I(",
                           rep(attr.vars, each = length(hvals)), "*",
                           rep(paste0("(",h, "=='", hvals, "')"), length(attr.vars)), ")")
      mf <-  as.formula(paste(o, " ~ ",  "factor(", h, ") + ", 
                              paste(c(restr.vars),
                                    collapse = " + "),
                              ifelse(!is.null(fsl),
                                     paste0(" | ", fsl),
                                     "")))
      
      ## Estimate
      mf <- fixest::feols(mf,
                          data = df,
                          cluster = clustvar)
      
      ## F-Test
      test <- wald(mf, 
                   keep = paste0(h, " =="), 
                   cluster = clustvar, print = F)
      
      ## Save for later
      omn.test.res <<- rbind(omn.test.res ,
                             data.frame(outcome =  names(outvars)[outvars == o],
                                        hetvar = h,
                                        wald_fstat = test$stat, 
                                        wald_pval = test$p))
      
      
      # Extract results
      res.df <- rbind(
        data.frame(outcome = names(outvars)[outvars == o],
                   outvar = o,
                   hetvar = h,
                   hetvar_val = as.numeric(rep(factor(hvals), length(amce.vars))),
                   hetname = het.names.vec[unlist(hetvars) == h],
                   hetvar_lab = rep(factor(hvals), length(amce.vars)),
                   attr = rep(as.numeric(gsub("\\_.*|attr", "", amce.vars)), each = length(hvals)),
                   level = rep(as.numeric(substr(amce.vars, nchar(amce.vars), nchar(amce.vars))), each = length(hvals)),
                   treat = rep(amce.vars, each = length(hvals)),
                   coef = m$coefficients[-seq_along(hvals)],
                   se = m$se[-seq_along(hvals)],
                   p = summary(m)$coeftable[-seq_along(hvals),4],
                   wald_stat = test$stat, 
                   wald_pval = test$p),
        data.frame(outcome = names(outvars)[outvars == o],
                   outvar = o,
                   hetvar = h,
                   hetname = het.names.vec[unlist(hetvars) == h],
                   hetvar_val = as.numeric(rep(factor(hvals), length(mm.vars[!mm.vars %in% amce.vars]))),
                   hetvar_lab = rep(factor(hvals), length(mm.vars[!mm.vars %in% amce.vars])),
                   attr =  rep(as.numeric(gsub("\\_.*|attr", "", mm.vars[!mm.vars %in% amce.vars])), each = length(hvals)),
                   level = rep(1, each = length(hvals)),
                   treat = rep(mm.vars[!mm.vars %in% amce.vars], each = length(hvals)),
                   coef = 0,
                   se = NA,
                   p = NA,
                   wald_stat = test$stat, 
                   wald_pval = test$p)
      )
      
      res.df$treat_lab = attr.lab.vec[res.df$treat]
      res.df$treat_type = names(attr.labs)[res.df$attr]
      res.df <- res.df[order(res.df$attr, res.df$level),]
      res.df
      
      ## Return
      return(res.df)
    }))
    
    # Combine all results for this variable and prepare plot
    coef.df <- prep_plot(coef.df)
    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)) / 2
    
    # print(coef.df)
    
    # Plot
    if(plot){
      het.name = het.names.vec[unlist(hetvars) == h]
      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(size = .85) +
        # geom_point(aes(y = yval + .15), col = "transparent") +
        # geom_point(aes(y = yval - .15), col = "transparent") +
        geom_segment(aes(yend = yval, 
                         x = coef - 1.96*se, 
                         xend = coef + 1.96*se)) +
        facet_grid(treat_type ~ outcome,
                   scales  = "free", space = "free_y",
                   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("Average Marginal Component Interaction Effect "),
             col = paste0(het.name,":"), 
             pch = paste0(het.name,":"),
             caption = paste0("Moderator variable: ", het.name, 
                              "\nRespondents: ",
                              paste(paste0(all.hvals, " (",sapply(all.hvals, function(x){sum(df[!duplicated(df$id), h] == x, na.rm = T)}), ")"),
                                    collapse = ", "),
                              "\nWald-test of nullity of sub-group differences in AMCEs:\n",
                              paste(sapply(unique(as.character(coef.df$outcome)), function(t){
                                paste0(t, " (F: ", 
                                       signif(coef.df$wald_stat[coef.df$outcome == t][1], 2), "; p: ",
                                       signif(coef.df$wald_pval[coef.df$outcome == t][1], 2), ")")
                              }), collapse = ", "))) +
        theme(panel.spacing = unit(.75, "lines"),
              strip.text.y.left = element_text(angle = 0, hjust = 0),
              legend.position = "top",
              panel.grid.major.y = element_blank(),
              panel.grid.minor.y = element_line(size = .5),
              legend.box.margin=margin(-10,-10,-10,-10),
              plot.caption = element_text(hjust = 0, colour = "grey20")) +
        scale_color_manual(values = custom.col)
      
      # Save
      png(file.path(fig.path, paste0(prefix, "_hetamce_",h,".png")),
          width = 7, height = 5, units = "in", res = 400)
      print(g)
      dev.off()
    } else {
      all.plot.df <- rbind(all.plot.df, coef.df)
    }
    
  }
  
  
  # OVERVIEW TABLE #################
  if(summary){
    # Adjust p-alue / Bonferroni
    omn.test.res$p.adjust <- stats::p.adjust(omn.test.res$wald_pval,
                                             method = "bonferroni")
    
    # Outcome factor
    omn.test.res$outcome <- factor(omn.test.res$outcome,
                                   levels = unique(omn.test.res$outcome),
                                   ordered = T)
    
    # Prepare table 
    
    ## Long -> wide by outcome
    tab <- do.call(cbind, split(omn.test.res[,c("wald_fstat", "wald_pval", "p.adjust")], 
                                omn.test.res$outcome))
    tab <- cbind(var = het.names.vec,
                 tab)
    
    ## Feed to stargazer
    tab.tex <- stargazer(tab,
                         title = "Omnibus Wald-Test Result for Joint Nullity of Heterogenous Effects by Moderator",
                         summary = FALSE,
                         column.sep.width = "0pt",
                         label = paste0("tab:", prefix, "heter_test"),
                         rownames = FALSE,
                         digits = 2, align = T,
                         notes.append = F, 
                         notes = paste0("Note: Adjusted p-values based on a Bonferroni adjustment for ", length(omn.test.res$wald_pval)," hypotheses."))
    
    ## Fix stargazer output
    tab.tex[11] <- paste0(paste(paste0(" \\multicolumn{1}{c}{",
                                       c("Moderator", rep(c("F-Stat", "$p$", "Adj. $p$"), 2)), "}"), 
                                collapse = "&"), "\\\\")
    tab.tex[9] <- paste(tab.tex[9], 
                        "& \\multicolumn{3}{c}{Score (0-1)} & \\multicolumn{3}{c}{Forced choice (0/1)} \\\\ \\\\[-1.8ex]")
    tab.tex[c(13 + c(0, cumsum(sapply(hetvars, length)[-length(hetvars)])))] <- paste0("\\\\[-1.8ex] \\multicolumn{1}{l}{\\textbf{",
                                                                                       names(hetvars),
                                                                                       "}} \\\\",
                                                                                       tab.tex[c(13 + c(0, cumsum(sapply(hetvars, length)[-length(hetvars)])))])
    tab.tex[13:length(tab.tex)] <- gsub("multicolumn{1}{c}", "multicolumn{1}{l}", 
                                        tab.tex[13:length(tab.tex)], fixed = T)
    
    # Write table to disc
    writeLines(tab.tex, 
               file.path(tab.path, paste0(prefix, "_het_test.tex")))
  }
  
  if(!plot){
    return(all.plot.df)
  }
  
}




# RANKING FUNCTIONS ################

get_order <- function(df, vars, 
                      get_coefs = FALSE){
  # Locals to keep track of things
  drop <- c()
  order <- c()
  pred <- rep(NA, nrow(df))
  drop.pairs <- c()
  eval.mat <- NULL
  models <- list()
  
  # Loop over all features (vars) until none left
  while(length(vars) > 0){
    
    # Estimate marginal means
    model.ls <- lapply(1:5, function(a){
      these.vars <- vars[grepl(paste0("attr", a), vars)]
      if(length(these.vars) <= 1){
        return(NULL)
      }
      # Formula
      f <- as.formula(paste("choice", " ~ 0 + ", 
                            paste(these.vars,
                                  collapse = " + "),
                            ""))
      
      this.df <-  df[!df$resp_pair %in% drop.pairs &
                              df[, paste0("attr", a, "_amb")] == 0,]
      
      # Estimate
      m <- try(fixest::feols(f,
                         data = this.df,
                         cluster = clustvar),
               silent = TRUE)
      if( "try-error" %in% class(m)){
        return(NULL)
      } else {
        return(m)
      }
    })
    names(model.ls) <- paste0("attr", 1:5)
    
    # Save model is so wished
    if(get_coefs){
      models <- c(models, 
                  list(model.ls))
    }
    
    # Stop if no more variation
    if(all(sapply(model.ls, is.null))){
      add.to.order <- sample(vars, size = length(vars), replace = F)
      add.to.order <- add.to.order[duplicated(gsub("d1|d2|d3", "amb", add.to.order))]
      order <- c(order, 
                 add.to.order)
      break
    }
    
    
    # Which is the next ranked feature to be dropped in next round?
    coefs <- lapply(model.ls, function(m){
      if(is.null(m)){
        c()
      } else {
        m$coefficients
      }
    })
    names(coefs) <- NULL
    coefs <- unlist(coefs)
    new.drop <- names(coefs)[which.max(abs(coefs-.5))]
    drop.a <- as.numeric(substr(new.drop, 5,5))
    if(length(model.ls[[drop.a]]$coefficients) == 2){
      new.drop <- c(new.drop, 
                    names(model.ls[[drop.a]]$coefficients)[
                      names(model.ls[[drop.a]]$coefficients) != new.drop
                    ])
    }
    order <- c(order, new.drop[1])
    drop <- c(drop, 
              new.drop) 
    amb <- gsub("d1|d2|d3", "amb", drop)
    
    # Drop variables
    vars <- vars[!vars %in% drop]
    
    # Drop data
    loss.obs <- apply(df[, drop, drop = F] == 1 & 
                        df[, amb, drop = F] == 0, 1, 
                      any)
    drop.pairs <- df$resp_pair[loss.obs]
  }
  if(get_coefs){
    return(models)
  } else {
    return(order)
  }
  
  
}

# Simple resampler based on respondent ID
resample <- function(df, resp.ids){
  new.ids <- sample(resp.ids, length(resp.ids), replace = T)
  new.df <- plyr::join(data.frame(id = new.ids),
                       df, by = "id", match = "all", type = "left")
  new.df
}


# Get rank at which strategy is evaluated 
decision_rank <- function(df, order){
  eval.mat <- NULL
  drop = c()
  rank <- rep(NA, nrow(df))
  for(v in order){
    # Which is the next one to drop
    drop <- c(drop, 
              v) 
    amb <- gsub("d1|d2|d3", "amb", drop)
    
    # Drop data
    loss.obs <- apply(df[, drop, drop = F] == 1 & 
                        df[, amb, drop = F] == 0, 1, 
                      any)
    drop.pairs <- df$resp_pair[loss.obs]
    
    # Save rank of dropped ones
    these.obs <- df$resp_pair %in% drop.pairs  & is.na(rank)
    rank[these.obs] <- which(order == v)
  }
  return(rank)
}

# Get F-Stats of contrasts between subsequent ranks
get_rank_fstats <- function(df, order, outcome =  "choice" ){
  ## Get rank at which obs is decided
  df$rank <- decision_rank(df, order)
  
  ## Make contrasts
  test.df <- do.call(rbind, lapply(1:(length(order)-1), function(i){
    ## Model 
    attr.vars <- order
    restr.vars <- paste0("I(",
                         rep(order[(i+1):length(order)], each = 1), "*",
                         rep(paste0("(rank >", i, ")"), 1), ")")
    mf <-  as.formula(paste(outcome, " ~ ", paste0("I(rank >", i, ") +"), 
                            paste(c(attr.vars, 
                                    restr.vars),
                                  collapse = " + ")))
    mf <- fixest::feols(mf,
                        data = df[df$rank >= i,],
                        cluster = clustvar)
    if(is.null(mf$collin.var)){
      test <- fixest::wald(mf, 
                           keep = paste0("rank >"), 
                           cluster = clustvar, print = FALSE)
    } else {
      test <- list(stat = NA, p = NA)
    }
    
    data.frame(rank = i, 
               N = mf$nobs,
               var = order[i],
               fstat = test$stat,
               pval = test$p
    )
  }))
  return(test.df)
}

# Get marginal mean by rank (nested marginal means)
get_nested_margmean <- function(df, order, outcome = "choice"){
  ## Get rank at which obs is decided
  df$rank <- decision_rank(df, order)
  
  ## Make contrasts
  do.call(rbind, lapply(1:(length(order)-1), function(i){
    do.call(rbind, lapply(1:5, function(a){
      # Formula
      attr <- paste0("attr",a, "_d", 1:3)
      attr <- attr[!attr %in% order[seq_len(i-1)]]
      if(length(attr) <= 1){
        return(NULL)
      }
      f <- as.formula(paste(outcome, " ~ 0 + ", 
                            paste(attr,
                                  collapse = " + ")))
      # Estimate
      m <- fixest::feols(f,
                         data = df[df[,paste0("attr", a, "_amb")] == 0 &
                                            df$rank >= i, ],
                         cluster = clustvar)
      
      # Extract results
      data.frame(outcome = names(outvars)[outvars == "choice"],
                 N = m$nobs,
                 rank = i, 
                 rankvar = c("Baseline", order)[i],
                 level = as.numeric(substr(attr, nchar(attr), nchar(attr))),
                 treat = attr,
                 treat_lab = as.vector(attr.labs[[a]])[as.numeric(substr(attr, nchar(attr), nchar(attr)))],
                 treat_type = names(attr.labs)[a],
                 coef = m$coefficients,
                 se = m$se,
                 p = summary(m)$coeftable[,4])
    }))
  }))
}

## Marginal means within ranks
get_within_rank_margmean <- function(df, order){
  ## Get rank at which obs is decided
  df$rank <- decision_rank(df, order)
  
  ## Make contrasts
  do.call(rbind, lapply(1:(length(order)-1), function(i){
    do.call(rbind, lapply(1:5, function(a){
      
      # Formula
      attr <- paste0("attr",a, "_d", 1:3)
      attr <- attr[!attr %in% order[seq_len(i-1)]]
      if(length(attr) <= 1){
        return(NULL)
      }
      f <- as.formula(paste("choice", " ~ 0 + ", 
                            paste(attr,
                                  collapse = " + ")))
      # Estimate
      m <- fixest::feols(f,
                         data = df[df[,paste0("attr", a, "_amb")] == 0 &
                                            df$rank == i, ],
                         cluster = clustvar)
      
      # Extract results
      data.frame(outcome = names(outvars)[outvars == "choice"],
                 N = m$nobs,
                 N_total = sum(df$rank == i),
                 rank = i, 
                 rankvar = c("Baseline", order)[i],
                 level = as.numeric(substr(attr, nchar(attr), nchar(attr))),
                 treat = attr,
                 treat_lab = as.vector(attr.labs[[a]])[as.numeric(substr(attr, nchar(attr), nchar(attr)))],
                 treat_type = names(attr.labs)[a],
                 coef = m$coefficients,
                 se = m$se,
                 p = summary(m)$coeftable[,4])
    }))
  }))
}

## Get ranking from order
get_ranks <- function(vars, order){
  sapply(vars, function(v){
    x <- which(order == v)
    if(length(x) == 0){
      x <- which(grepl(substr(v, 1, nchar(v)-1), order, fixed = T))[-1]
    }
    x
  })
}

