library(ggplot2)
library(dplyr)
library(tidyr)
library(tibble)
library(readr)
library(forcats)
library(knitr)
library(kableExtra)

set.seed(215)

source("../../admin_scripts/util.R")

ballot <- read.csv("../../data/common_pleas/judicial_ballot_position.csv")
ballot$name <- tolower(ballot$name)
ballot$name <- gsub("[[:punct:]]", " ", ballot$name)
ballot$name <- trimws(ballot$name)
ballot$year <- as.character(ballot$year)

df_all <- safe_load("C:/Users/Jonathan Tannen/Dropbox/sixty_six/data/processed_data/df_all_2019_07_09.Rda")

df <- df_all %>% 
  filter(
    election == "primary",
    grepl("COMMON PLEAS$", OFFICE),
    PARTY == "DEMOCRATIC",
    CANDIDATE != "Write In",
    year >= 2009
  )
rm(df_all)

df <- df %>%
  mutate(
      CANDIDATE = tolower(CANDIDATE),
      CANDIDATE = gsub("\\s+", " ", CANDIDATE)
    ) %>%
  group_by(CANDIDATE, year, WARD19, DIV19) %>%
  summarise(votes = sum(VOTES)) %>%
  ungroup()

df_ward <- df %>%
  group_by(CANDIDATE, year, WARD19) %>%
  summarise(votes = sum(votes)) %>%
  group_by(year, WARD19) %>%
  mutate(pvote = votes / sum(votes)) %>%
  ungroup

df_total <- df_ward %>% 
  group_by(year, CANDIDATE) %>% 
  summarise(votes = sum(votes))

election <- tibble(
  year = as.character(c(2009, 2011, 2013, 2015, 2017, 2019)),
  votefor = c(7, 10, 6, 12, 9, 6)
)

election <- election %>% left_join(
  ballot %>% 
    mutate(year = as.character(year)) %>%
    group_by(year) %>% 
    summarise(
      nrows = max(rownumber),
      ncols = max(colnumber), 
      ncand = n(),
      n_philacomm = sum(philacommrec),
      n_inq = sum(inq),
      n_dcc = sum(dcc)
    )
)

df_total <- df_total %>% 
  inner_join(election) %>%
  group_by(year) %>%
  arrange(desc(year), desc(votes)) %>%
  mutate(finish = 1:n()) %>%
  mutate(winner = finish <= votefor)

df_total <- df_total %>% inner_join(
  ballot,
  by = c("CANDIDATE" = "name", "year" = "year")
)

df_total <- df_total %>%
  group_by(year) %>%
  mutate(pvote = votes / sum(votes))

df <- df %>%
  inner_join(ballot, by = c("CANDIDATE"="name", "year"="year"))

df_ward <- df_ward %>%
  inner_join(ballot, by = c("CANDIDATE"="name", "year"="year"))


prep_data <- function(df){
  df %>% mutate(
    rownumber = fct_relevel(factor(as.character(rownumber)), "3"),
    colnumber = fct_relevel(factor(as.character(colnumber)), "3"),
    col1 = colnumber == 1,
    col2 = colnumber == 2,
    col3 = colnumber == 3,
    row1 = rownumber == 1,
    row2 = rownumber == 2,
    col1row1 = (colnumber==1) & (rownumber==1), 
    col1row2 = (colnumber==1) & (rownumber==2), 
    is_rec = philacommrec > 0,
    is_highly_rec = philacommrec==2,
    inq=inq>0,
    candidate_year = paste0(CANDIDATE, "::", year)
  )
}

df_ward <- prep_data(df_ward)

df_ward <- df_ward %>%
  group_by(year, WARD19) %>%
  mutate(pvote = votes / sum(votes)) %>% 
  ungroup

library(lme4)

# better opt: https://github.com/lme4/lme4/issues/98
library(nloptr)
defaultControl <- list(
  algorithm="NLOPT_LN_BOBYQA",xtol_rel=1e-6,maxeval=1e5
)
nloptwrap2 <- function(fn,par,lower,upper,control=list(),...) {
    for (n in names(defaultControl))
      if (is.null(control[[n]])) control[[n]] <- defaultControl[[n]]
    res <- nloptr(x0=par,eval_f=fn,lb=lower,ub=upper,opts=control,...)
    with(res,list(par=solution,
                  fval=objective,
                  feval=iterations,
                  conv=if (status>0) 0 else status,
                  message=message))
}

USE_SAVED <- TRUE
rfit_filename <- "outputs/lmer_fit.Rds"
if(!USE_SAVED){
  rfit <- lmer(
    log(pvote + 0.001) ~ 
      (1 | candidate_year)+
      col1row1 +
      col1row2+
      row1 + row2 +
      I(gender == "F") +
      col1 + col2 +
      dcc + 
      is_rec + is_highly_rec +
      factor(year) +
      (
        # row1 + row2 +
        col1row1 +
          col1row2 +
          I(gender == "F") +
          col1 + col2 + #col3 +
          dcc +
          is_rec + is_highly_rec
        | WARD19
      ),
    df_ward,
    control=lmerControl(optimizer="nloptwrap2")
  )
  saveRDS(rfit, file=rfit_filename)
} else {
  rfit <- readRDS(rfit_filename)
}

df_ward$fitted <- predict(rfit)

ranef <- as.data.frame(ranef(rfit)$WARD19) %>% 
  rownames_to_column("WARD19")  %>%
  gather("variable", "random_effect", -WARD19) %>%
  mutate(
    fixed_effect = fixef(rfit)[variable],
    effect = random_effect + fixed_effect
  )

ranef_wide <- ranef %>% 
  select(WARD19, variable, effect) %>% 
  spread(variable, effect)

df_pred <- df %>% 
  group_by(year, WARD19, DIV19) %>%
  mutate(pvote = votes/sum(votes)) %>% 
  group_by() %>%
  left_join(ranef_wide, by="WARD19") %>%
  mutate(
    log_pvote = log(pvote + 0.001)
  ) %>%
  mutate(
    log_pvote_adj = log_pvote -
      ifelse(
        colnumber == 1, col1TRUE, 0 
      ) - 
      ifelse(
        colnumber == 2, col2TRUE, 0 
      ) - 
      ifelse(
        colnumber == 1 & rownumber == 1, col1row1TRUE, 0
      ) - 
      ifelse(
        colnumber == 1 & rownumber == 2, col1row2TRUE, 0
      )
  )

turnout <- df_pred %>% group_by(year, WARD19, DIV19) %>%
  summarise(votes=sum(votes)) %>%
  ungroup()

run_elections <- function(
  df, year, order, nrow=11, ncol=3, n_elec=1e3
){
  
  if(!all(c("CANDIDATE", "order") %in% colnames(order))){
    stop(
      "order should be a dataframe with columns CANDIDATE, order, and maybe WARD19 DIV19"
    )
  }
  
  order <- order %>% mutate(
    col = ((order-1) %/% nrow)+1,
    row = order - (col-1)*nrow
  )
  
  ## assume everything is correctly sorted

  df$log_pvote <- df %>% 
    left_join(order) %>%
    with(
      log_pvote_adj + 
        ifelse(col==1, col1TRUE, 0) +
        ifelse(col==2, col2TRUE, 0) +
        ifelse(col==1 & row==1, col1row1TRUE, 0) +
        ifelse(col==1 & row==2, col1row2TRUE, 0)
    )
  
  df$pvote <- exp(df$log_pvote)
  
  votes <- df %>%
    ungroup() %>%
    select(CANDIDATE, WARD19, DIV19, pvote, turnout, year) %>%
    mutate(expected_votes = pvote * turnout) %>%
    group_by(CANDIDATE, year) %>%
    summarise(expected_votes = sum(expected_votes)) %>%
    ungroup() %>%
    left_join(
      expand.grid(year=as.character(year), elec=1:n_elec)
    ) %>%
    mutate(
      votes = rpois(n(), lambda=expected_votes)
    )
  
  return(votes %>% select(CANDIDATE, elec, votes))
}

run_sims <- function(
  get_order_fn, year, df=df_pred, nsims=400, 
  nelec=1e2, verbose=TRUE, 
  return_orders=TRUE, return_results=TRUE
){
  df <- df %>% 
    filter(year == !!year) %>%
    arrange(WARD19, DIV19, CANDIDATE)
  df <- df %>%
    left_join(
      turnout %>% filter(year == !!year) %>% rename(turnout=votes)
    )
  results <- vector(mode="list", length=nsims)
  orders <- vector(mode="list", length=nsims)
  win_fracs <- vector(mode="list", length=nsims)
  
  for(i in 1:nsims){
    order_sim <- get_order_fn(df) %>% 
      mutate(sim=i)

    if(return_orders) orders[[i]] <- order_sim
    
    result <- run_elections(
      df, year, order_sim, n_elec=nelec
    ) %>% mutate(sim=i)
    
    if(return_results) results[[i]] <- result
    
    win_fracs[[i]] <- get_win_frac(result, year)
    if(verbose) if(i %% 10 == 0) print(i)

  }
  
  return_list <- list(
    win_fracs=bind_rows(win_fracs), 
    year=year
  )

  return_list$dev <- get_dev(return_list$win_fracs)
  
  if(return_results) return_list$results <- bind_rows(results)
  if(return_orders) return_list$orders <- bind_rows(orders)
  
  return(return_list)
}

get_win_frac <- function(
  results,
  year
){
  nwinners <- election$votefor[election$year == year]
  
  results <- results %>%
    group_by(sim, elec) %>%
    mutate(
      pvote = votes/sum(votes),
      rank = rank(desc(votes), ties.method="first")
    ) %>%
    ungroup() %>%
    mutate(is_winner = (rank <= nwinners)) %>%
    group_by(sim, CANDIDATE) %>%
    summarise(
      pwin = mean(is_winner),
      pvote = mean(pvote)
    )

  return(results)
}


get_dev <- function(frac_winners){
  frac_winners <- frac_winners %>%
    group_by(CANDIDATE) %>%
    summarise(
      dev = mean(abs(pwin - mean(pwin))),
      pwin = mean(pwin)
    )
  
  return(frac_winners)
}

ggplot(
  df_total %>% 
    filter(year == 2017) %>% 
    mutate(
      lastname=format_name(gsub(".*\\s(\\S+)$", "\\1", CANDIDATE))
    ) %>%
    arrange(votes),
  aes(y=rownumber, x=colnumber)
) +
  geom_tile(
    aes(fill=pvote*100, color=winner),
    size=2
  ) +
  geom_text(
    aes(
      label = ifelse(philacommrec==1, "R", ifelse(philacommrec==2,"HR","")),
      x=colnumber+0.45,
      y=rownumber+0.45
    ),
    color="grey70",
    hjust=1, vjust=0
  ) +
  geom_text(
    aes(
      label = ifelse(dcc==1, "D", ""),
      x=colnumber-0.45,
      y=rownumber+0.45
    ),
    color="grey70",
    hjust=0, vjust=0
  ) +
  geom_text(
    aes(label = sprintf("%s\n%0.1f%%", lastname, 100*pvote)),
    color="black"
    # fontface="bold"
  ) +
  scale_y_reverse(NULL) +
  scale_x_continuous(NULL)+
  scale_fill_viridis_c(guide=FALSE) +
  scale_color_manual(values=c(`FALSE`=NA, `TRUE`="yellow"), guide=FALSE) +
  annotate(
    "text",
    label="R = Recommended\nD = DCC Endorsed",
    x = 1.6,
    y = 11,
    hjust=0,
    color="grey70"
  ) +
  expand_limits(x=3.5)+
  theme_sixtysix() %+replace% 
  theme(
    panel.grid.major=element_blank(),
    axis.text=element_blank()
  ) +
  ggtitle(
    "2017 Common Pleas Results",
    "Democratic Primary. Winners are outlined."
  )

cali_ward_order <- function(df){
  candidates <- df %>%
    with(sample(unique(CANDIDATE)))
  
  wards <- sprintf("%02d", 1:66)

  cand_order <- data.frame()
  
  for(ward in wards){
      cand_order <- bind_rows(
        cand_order,
        data.frame(
          WARD19=ward,
          CANDIDATE=candidates,
          order=seq_along(candidates)
        )
      )
      last <- length(candidates)
      candidates <- c(candidates[last], candidates[-last])
  }
  return(cand_order)
}

cali_ward_results <- run_sims(
  cali_ward_order, 2017, nsims=1e2, nelec=1e2, verbose=FALSE
)

summary_table(cali_ward_results)

# summary_stats(cali_ward_results)

methods <- c(
  status_quo = status_quo_order,
  random_div = random_div_order,
  random_ward = random_ward_order,
  cycle_ward = cycle_ward_order,
  cali_ward = cali_ward_order
)

results_df <- data.frame()
VERBOSE <- TRUE
USE_SAVED <- TRUE
SAVE_RESULTS <- TRUE

for(year in seq(2009, 2019, 2)){
  if(VERBOSE) print(year)
  for(method in names(methods)){
    if(VERBOSE)  print(method)
    filename <- sprintf("%s_%s.Rds", method, year)
    filedir <- "outputs"
    
    if(!USE_SAVED | !filename %in% list.files(filedir)){
      sim <- run_sims(
        methods[[method]], year, 
        verbose=VERBOSE, return_orders=FALSE,
        return_results=FALSE
      )
      summ_stats <- summary_stats(sim)
      
      if(SAVE_RESULTS) {
        saveRDS(summ_stats, file = paste0(filedir, "/", filename))
      }
    } else {
      summ_stats <- readRDS(paste0(filedir, "/", filename))
    }
    
    results_df <- bind_rows(
      results_df,
      data.frame(
        summ_stats,
        year=year,
        method=method
      )
    )
  }
}

replace_names <- c(
  dev = "Average Absolute Deviation (%)",
  n_rec = "Recommended Winners",
  n_dcc = "DCC-Endorsed Winners"
)

for(i in 1:length(replace_names)){
  names(results_df)[
    which(names(results_df) == names(replace_names)[i])
    ] <- replace_names[i]
}

results_df$`Average Absolute Deviation (%)` <- results_df$`Average Absolute Deviation (%)`*100

replace_method <- c(
  status_quo = "Status Quo",
  random_div = "Random by Division",
  random_ward = "Random by Ward",
  cycle_ward = "Wheel Cycle by Ward", 
  cali_ward = "California Cycle by Ward"
)

results_df$method <- replace_method[results_df$method]

out_df <- results_df %>% 
    gather_("key", "value", replace_names) %>%
    unite(key, key, method, sep=":\n") %>%
    spread(key, value)

col_order <- expand.grid(a=replace_method, b=replace_names) %>%
  with(paste0(b, ":\n", a))

out_df <- df_total %>% 
  group_by(year) %>%
  summarise(
    `Number of\nWinners` = votefor[1],
    `Observed\nRecommended Winners` = sum(winner * (philacommrec > 0)),
    `Observed\nDCC-Endorsed Winners` = sum(winner * dcc)
  ) %>%
  mutate(year = asnum(year)) %>%
  left_join(out_df)

names(out_df)[1] <- "Year"

avgs <- results_df %>%
  group_by(method) %>%
  summarise_all(funs(sprintf("%0.1f", mean(.))))

status_quo_dev <- avgs$"Average Absolute Deviation (%)"[avgs$method == "Status Quo"]
random_div_dev <- avgs$"Average Absolute Deviation (%)"[avgs$method == "Random by Division"]
random_ward_dev <- avgs$"Average Absolute Deviation (%)"[avgs$method == "Random by Ward"]
wheel_ward_dev <- avgs$"Average Absolute Deviation (%)"[avgs$method == "Wheel Cycle by Ward"]


kable(
  out_df[,c(names(out_df)[1:4], col_order)],
  digits=c(rep(0,4), rep(1,length(col_order))),
  col.names=c(names(out_df)[1:4], gsub(".*:\\n(.*)", "\\1", col_order)),
    padding = "4in"
) %>% 
  add_header_above(c(
    " "=4,
    "Average Absolute Deviation (%)"=length(replace_method),
    "Recommended Winners"=length(replace_method),
    "DCC-Endorsed Winners"=length(replace_method)
  )) %>%
  kable_styling() %>%
  scroll_box(width="100%")

plot_df <- results_df %>% left_join(
  out_df %>% rename(year=Year, nwin = `Number of\nWinners`) %>%
    select(year, nwin)
) %>%
  mutate(
    `% Recommended` = `Recommended Winners`/nwin * 100,
    `% DCC Endorsed` = `DCC-Endorsed Winners`/nwin * 100
  )

cols <- c(
  "Average Absolute Deviation (%)",
  "% Recommended", 
  "% DCC Endorsed"
)

plot_var <- function(var, title = NULL, subtitle=NULL){
  var <- enquo(var)
  ggplot(
    plot_df,
    aes(
      y = !!var, 
      x = factor(method, levels = replace_method),
      group=interaction(method, year), fill=as.character(year)
    )
  ) +
    geom_bar(stat="identity", position="dodge", color="black") +
    xlab(NULL) + ylab(rlang::as_name(var)) +
    scale_fill_discrete(
      NULL, 
      h = c(253, 253), c = c(50, 100), 
      l = seq(10, 80, length.out = 6)
    ) +
    theme_sixtysix() %+replace%
    theme(
      axis.text.x = element_text(angle=90, hjust=1),
      legend.position = "right",
      legend.direction = "vertical"
    ) +
    ggtitle(title, subtitle)
}

plot_var(
  `Average Absolute Deviation (%)`, 
  "All methods drastically reduce the effect of the lottery"
)

plot_var(
  `% Recommended`, 
  "Eliminating the Status Quo slightly the Recommended judges",
  "Percentage increases represent about 0.5 more Rec. judges per year"
)

plot_var(
  `% DCC Endorsed`, 
  "Eliminating the Status Quo increases the DCC's effect",
   "Percentage increases represent about 0.6 more DCC endorsed judges per year"
)