# Script to create preference flow estimates from the distribution of preferences
# (all years), and compare the results to the actual flows (1996-2013).

library(ggplot2)
library(readr)
library(dplyr)


setwd("flow_estimates/")

# For the second run-through: do we want to use the
# minimal-preference data set for testing the TPP
# estimates?
use_minimal_data = 0

in_folder = "../exact_data"

is_classic = function(parties) {
  # Function only called during a 1996-2013 loop,
  # so don't need the full set of Coalition parties.
  coalition_parties = c("Lib46", "NPA", "CLP", "LNP")
  
  classic = FALSE
  
  if ("ALP" %in% parties) {
    ALP_i = which(parties == "ALP")
    OTH_i = 3 - ALP_i
    
    if (parties[OTH_i] %in% coalition_parties) {
      classic = TRUE
    }
  }
  
  return(classic)
}

png_print = function(gg_img, out_file, width=500, height=500) {
  png(filename=out_file, width=width, height=height, units="px")
  print(gg_img)
  dev.off()
  return(0)
}

estimate_flows = function(seat_prefs.df) {
  # Input a seat_prefs data frame.
  # *** Think about cand_party fields ***
  # Outputs a list with results for standard and simple estimated preference flows.
  
  last_count = max(seat_prefs.df$Count)
  count.df = filter(seat_prefs.df, Count == last_count)
  
  if ("Exhausted" %in% count.df$To_candidate) {
    # Remove the exhausted line
    count.df = count.df[-which(count.df$To_candidate == "Exhausted"), ]
  }
  
  if (length(count.df$Seat) == 2) {
    # Full distribution of preferences
    
    cand_party1 = count.df$To_cand_party[1]
    cand_party2 = count.df$To_cand_party[2]
    
    tcp_cands = c(count.df$To_candidate[1], count.df$To_candidate[2])
    tcp_parties = c(count.df$To_party[1], count.df$To_party[2])
    tcp_cand_parties = c(cand_party1, cand_party2)
    
    # In the flow_rates list, the flows vector gives flows on a scale of 0-1
    # for:
    # c(cand1, cand2, exhausted)
    
    flow_rates = list()
    flow_rates_cand_parties = c(cand_party1, cand_party2, "Exhausted_NA")
    
    flow_rates_cands = c(tcp_cands[1], tcp_cands[2], "Exhausted")
    flow_rates_parties = c(tcp_parties[1], tcp_parties[2], "Exhausted")
    
    
    simple_flow_rates = list()
    simple_list_entries = 0
    
    flow_rates[[1]] = c(1, 0, 0)
    flow_rates[[2]] = c(0, 1, 0)
    flow_rates[[3]] = c(0, 0, 1)
    list_entries = 3
    
    # Work in reverse order through the preference counts.
    counts = unique(seat_prefs.df$Count)
    counts = counts[order(-counts)]
    
    for (count in counts) {
      count.df = filter(seat_prefs.df, Count == count)
      
      # Need to do this because we want to properly keep track of the
      # exhausted fraction (default percentages add up to > 100% because
      # exhausted votes aren't included in the Transferred_perc to candidates).
      sum_votes = sum(count.df$Transferred_votes)
      count.df$Real_transf_perc = count.df$Transferred_votes / sum_votes
      
      this_from_cand_party = count.df$From_cand_party[1]
      this_flow = c(0, 0, 0)
      
      for (i in 1:length(count.df$Count)) {
        this_to_cand_party = count.df$To_cand_party[i]
        this_i = which(flow_rates_cand_parties == this_to_cand_party)
        
        if (length(this_i) != 1) {
          print("Couldn't find to_cand_party")
          stop()
        }
        
        this_perc = count.df$Real_transf_perc[i]
        this_flow = this_flow + this_perc*flow_rates[[this_i]]
      }
      
      count_tcp_cands.df = filter(count.df, To_cand_party %in% tcp_cand_parties)
      sum_tcp_transf_votes = sum(count_tcp_cands.df$Transferred_votes)
      
      first_i = which(count_tcp_cands.df$To_cand_party == tcp_cand_parties[1])
      simple_flow1 = round(100 * count_tcp_cands.df$Transferred_votes[first_i] / sum_tcp_transf_votes, 2)
      simple_flow2 = 100 - simple_flow1
      
      flow_rates_cand_parties = c(flow_rates_cand_parties, this_from_cand_party)
      flow_rates_cands = c(flow_rates_cands, count.df$From_candidate[1])
      flow_rates_parties = c(flow_rates_parties, count.df$From_party[1])
      
      list_entries = list_entries + 1
      flow_rates[[list_entries]] = this_flow
      
      simple_list_entries = simple_list_entries + 1
      simple_flow_rates[[simple_list_entries]] = c(simple_flow1, simple_flow2)
    }
    
    combined_list = list(cand_parties = flow_rates_cand_parties,
                         from_cands = flow_rates_cands,
                         from_parties = flow_rates_parties,
                         standard_ests = flow_rates,
                         simple_ests = simple_flow_rates)
  } else {
    # Not a full distribution of preferences.
    combined_list = list()
  }
  
  return(combined_list)
}


# First run through: test against known flows, 1996-2013.  Outputs flow CSV files with
# comparisons, and PNG's of true v estimated scatter plots.

years = c(1996, 1998, 2001, 2004, 2007, 2010, 2013)

for (year in years) {
  infile_prim = sprintf("%s/prim_%d.csv", in_folder, year)
  infile_prefs = sprintf("%s/prefs_%d.csv", in_folder, year)
  infile_flow = sprintf("%s/flow_%d.csv", in_folder, year)
  
  outfile_flow = sprintf("est_flow_%d.csv", year)
  
  prim.df = read_csv(infile_prim)
  prefs.df = read_csv(infile_prefs)
  flow.df = read_csv(infile_flow)
  
  flow.df$Est_perc1 = 0
  flow.df$Est_perc2 = 0
  flow.df$Est_exh = 0
  
  flow.df$Simple_est_perc1 = 0
  flow.df$Simple_est_perc2 = 0
  
  # Concatenating these to avoid conflicts with the two Paul Keatings in Blaxland 1987:
  prim.df$Cand_party = sprintf("%s_%s", prim.df$Candidate, prim.df$Party)
  prefs.df$From_cand_party = sprintf("%s_%s", prefs.df$From_candidate, prefs.df$From_party)
  prefs.df$To_cand_party = sprintf("%s_%s", prefs.df$To_candidate, prefs.df$To_party)
  flow.df$From_cand_party = sprintf("%s_%s", flow.df$From_candidate, flow.df$From_party)
  flow.df$To_cand_party1 = sprintf("%s_%s", flow.df$To_candidate1, flow.df$To_party1)
  flow.df$To_cand_party2 = sprintf("%s_%s", flow.df$To_candidate2, flow.df$To_party2)
  
  seats = unique(flow.df$Seat)
  
  for (seat in seats) {
    seat_prim.df = filter(prim.df, Seat == seat)
    seat_prefs.df = filter(prefs.df, Seat == seat)
    seat_flow_i = which(flow.df$Seat == seat)
    seat_flow.df = flow.df[seat_flow_i, ]
    
    if (length(seat_prim.df$Seat) > 2) {
      last_count = max(seat_prefs.df$Count)
      count.df = filter(seat_prefs.df, Count == last_count)
      
      if ("Exhausted" %in% count.df$To_candidate) {
        # Remove the exhausted line
        count.df = count.df[-which(count.df$To_candidate == "Exhausted"), ]
      }
      
      if (is_classic(count.df$To_party)) {
        est_flows = estimate_flows(seat_prefs.df)
        flow_rates_cand_parties = est_flows$cand_parties
        flow_rates = est_flows$standard_ests
        simple_flow_rates = est_flows$simple_ests
        
        for (i in seat_flow_i) {
          j = which(flow_rates_cand_parties == flow.df$From_cand_party[i])
          if (length(j) != 1) {
            print("Couldn't find j")
            stop()
          }
          
          flow_orig = flow_rates[[j]]
          
          # Re-calculate the flows with the usual convention on exhausted votes.
          new_flow = flow_orig / (flow_orig[1] + flow_orig[2])
          new_flow[3] = flow_orig[3]
          
          new_flow = round(100 * new_flow, 2)
          
          if (flow.df$To_cand_party1[seat_flow_i[1]] == flow_rates_cand_parties[1]) {
            i1 = 1
            i2 = 2
          } else {
            i1 = 2
            i2 = 1
          }
          
          flow.df$Est_perc1[i] = new_flow[i1]
          flow.df$Est_perc2[i] = new_flow[i2]
          flow.df$Est_exh[i] = new_flow[3]
          
          flow.df$Simple_est_perc1[i] = simple_flow_rates[[j-3]][i1]
          flow.df$Simple_est_perc2[i] = simple_flow_rates[[j-3]][i2]
        }
      }
    }
  }
  
  flow.df = filter(flow.df, Est_perc1 != 0)
  flow.df$Error = flow.df$Est_perc1 - flow.df$Vote_perc1
  flow.df$Total_transf = flow.df$Votes1 + flow.df$Votes2
  
  this_scatter = ggplot() + geom_point(data=flow.df, aes(x=Est_perc1, y=Vote_perc1)) + xlab("True") + ylab("Estimated")
  this_simple_scatter = ggplot() + geom_point(data=flow.df, aes(x=Vote_perc1, y=Simple_est_perc1))
  
  outfile_img = sprintf("scatter_true_est_flow_%d.png", year)
  outfile_simple_img = sprintf("scatter_true_simple_est_flow_%d.png", year)
  
  png_print(this_scatter, outfile_img, 300, 300)
  png_print(this_simple_scatter, outfile_simple_img, 700, 700)
  
  write.csv(flow.df, file=outfile_flow, row.names=FALSE)
  
  r_1 = cor(flow.df$Vote_perc1, flow.df$Est_perc1)
  r_2 = cor(flow.df$Vote_perc1, flow.df$Simple_est_perc1)
  print(sprintf("%d: Standard method r = %.2f, simple method r = %.2f",
                year, r_1, r_2))
  
}


# Second run through: estimate flows 1919-1990

years = c(1919, 1922, 1925, 1928, 1929, 1931, 1934, 1937, 1940, 1943, 1946, 1949, 1951,
          1954, 1955, 1958, 1961, 1963, 1966, 1969, 1972, 1974, 1975, 1977, 1980, 1983,
          1984, 1987, 1990, 1993)

if (use_minimal_data == 1) {
  in_folder = "../tpp_estimates/minimal_pref_data"
  out_folder = "../tpp_estimates/minimal_pref_data"
  years = c(years, 1996, 1998, 2001, 2004, 2007, 2010, 2013)
} else {
  out_folder = "."
}

for (year in years) {
  print(year)
  infile_prim = sprintf("%s/prim_%d.csv", in_folder, year)
  infile_prefs = sprintf("%s/prefs_%d.csv", in_folder, year)
  
  prim.df = read_csv(infile_prim)
  prefs.df = read_csv(infile_prefs)
  flow.df = read_csv(infile_flow)
  
  # Concatenating these to avoid conflicts with the two Paul Keatings in Blaxland 1987:
  prim.df$Cand_party = sprintf("%s_%s", prim.df$Candidate, prim.df$Party)
  prefs.df$From_cand_party = sprintf("%s_%s", prefs.df$From_candidate, prefs.df$From_party)
  prefs.df$To_cand_party = sprintf("%s_%s", prefs.df$To_candidate, prefs.df$To_party)
  
  seats = unique(prefs.df$Seat)
  
  # "Seat","From_candidate","From_party","To_candidate1","To_party1","Votes1","Vote_perc1","To_candidate2","To_party2","Votes2","Vote_perc2","Exhausted","Exhausted_perc","Est_perc1","Est_perc2","Est_exh","Simple_est_perc1","Simple_est_perc2","From_cand_party","Error","Total_transf"
  Out_Seat = character(0)
  Out_From_candidate = character(0)
  Out_From_party = character(0)
  Out_To_candidate1 = character(0)
  Out_To_party1 = character(0)
  Out_Votes1 = numeric(0)
  Out_Vote_perc1 = numeric(0)
  Out_To_candidate2 = character(0)
  Out_To_party2 = character(0)
  Out_Votes2 = numeric(0)
  Out_Vote_perc2 = numeric(0)
  Out_Exhausted = numeric(0)
  Out_Exhausted_perc = numeric(0)
  Out_Calc_type = character(0)
  
  for (seat in seats) {
    seat_prim.df = filter(prim.df, Seat == seat)
    seat_prefs.df = filter(prefs.df, Seat == seat)
    
    est_flows = estimate_flows(seat_prefs.df)
    
    if (length(est_flows) > 0) {
      flow_rates_cand_parties = est_flows$cand_parties
      flow_rates = est_flows$standard_ests
      simple_flow_rates = est_flows$simple_ests
      
      last_count = max(seat_prefs.df$Count)
      count.df = filter(seat_prefs.df, Count == last_count)
      
      if ("Exhausted" %in% count.df$To_candidate) {
        # Remove the exhausted line
        count.df = count.df[-which(count.df$To_candidate == "Exhausted"), ]
      }
      
      
      if (length(seat_prim.df$Seat) == 3) {
        calc_type = "Exact"
      } else {
        calc_type = "Estimate"
      }
      
      est_flows = estimate_flows(seat_prefs.df)
      
      flow_rates_cand_parties = est_flows$cand_parties
      flow_rates_cands = est_flows$from_cands
      flow_rates_parties = est_flows$from_parties
      flow_rates = est_flows$standard_ests
      simple_flow_rates = est_flows$simple_ests
      
      # Loop starts at i=4 because i = {1, 2} are the TCP candidates,
      # and i = 3 is exhausted.
      for (i in 4:length(flow_rates)) {
        flow_orig = flow_rates[[i]]
        
        prim_i = which(seat_prim.df$Cand_party == flow_rates_cand_parties[i])
        if (length(prim_i) != 1) {
          print("Couldn't find primary votes")
          stop()
        }
        this_prim_votes = seat_prim.df$Votes[prim_i]
        
        # Order this weirdly to ensure that the votes tally:
        this_votes1 = round(flow_orig[1] * this_prim_votes)
        this_exh = round(flow_orig[3] * this_prim_votes)
        this_votes2 = this_prim_votes - this_votes1 - this_exh
        
        # Re-calculate the flows with the usual convention on exhausted votes.
        new_flow = flow_orig / (flow_orig[1] + flow_orig[2])
        new_flow[3] = flow_orig[3]
        
        new_flow = round(100*new_flow, 2)
        
        Out_Seat = c(Out_Seat, seat)
        Out_From_candidate = c(Out_From_candidate, flow_rates_cands[i])
        Out_From_party = c(Out_From_party, flow_rates_parties[i])
        Out_To_candidate1 = c(Out_To_candidate1, flow_rates_cands[1])
        Out_To_party1 = c(Out_To_party1, flow_rates_parties[1])
        Out_Votes1 = c(Out_Votes1, this_votes1)
        Out_Vote_perc1 = c(Out_Vote_perc1, new_flow[1])
        Out_To_candidate2 = c(Out_To_candidate2, flow_rates_cands[2])
        Out_To_party2 = c(Out_To_party2, flow_rates_parties[2])
        Out_Votes2 = c(Out_Votes2, this_votes2)
        Out_Vote_perc2 = c(Out_Vote_perc2, new_flow[2])
        Out_Exhausted = c(Out_Exhausted, this_exh)
        Out_Exhausted_perc = c(Out_Exhausted_perc, new_flow[3])
        Out_Calc_type = c(Out_Calc_type, calc_type)
      }
    }
  }
  
  out_flow.df = data.frame(Seat = Out_Seat,
                           From_candidate = Out_From_candidate,
                           From_party = Out_From_party,
                           To_candidate1 = Out_To_candidate1,
                           To_party1 = Out_To_party1,
                           Votes1 = Out_Votes1,
                           Vote_perc1 = Out_Vote_perc1,
                           To_candidate2 = Out_To_candidate2,
                           To_party2 = Out_To_party2,
                           Votes2 = Out_Votes2,
                           Vote_perc2 = Out_Vote_perc2,
                           Exhausted = Out_Exhausted,
                           Exhausted_perc = Out_Exhausted_perc,
                           Calc_type = Out_Calc_type,
                           stringsAsFactors = FALSE)
  
  outfile_flow = sprintf("%s/flow_%d.csv", out_folder, year)
  write.csv(out_flow.df, file=outfile_flow, row.names=FALSE)
}
