# Script to create the TPP estimates.  A few steps to get there.

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

options(warn=1)

# Set the following to 1 to use the minimal preference data for 1983-2013,
# for testing against the exact TPP values:
use_minimal_data = 1

setwd("tpp_estimates/")

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)

if (use_minimal_data == 0) {
  exact_folder = "../exact_data"
  flow_folder = "../flow_estimates"
} else {
  exact_folder = "minimal_pref_data"
  flow_folder = "minimal_pref_data"
  years = c(years, 1987, 1990, 1993, 1996, 1998, 2001, 2004, 2007, 2010, 2013)
}

source("aux_functions.R")

# Delete the files made last time this script was run:
prev_party_files = list.files(path=".", pattern="party_.*\\.csv")
file.remove(prev_party_files)

##############################################################
# Part 1: Create a CSV file for each party containing its
# (exact or estimated) preference flows.

for (i_year in seq_along(years)) {
  year = years[i_year]
  infile_flow = sprintf("%s/flow_%d.csv", flow_folder, year)
  in_flow.df = read_csv(infile_flow)
  
  for (i in seq_along(in_flow.df$From_candidate)) {
    tcp_parties = c(in_flow.df$To_party1[i], in_flow.df$To_party2[i])
    tpp_i = tpp_indices(tcp_parties)
    seat = in_flow.df$Seat[i]
    
    if (tpp_i[1] == 1) {
      # TPP contest
      if (tpp_i[2] == 1) {
        # ALP is party1
        ALP_votes = in_flow.df$Votes1[i]
        ALP_vote_perc = in_flow.df$Vote_perc1[i]
        
        LNP_votes = in_flow.df$Votes2[i]
        LNP_vote_perc = in_flow.df$Vote_perc2[i]
      } else {
        ALP_votes = in_flow.df$Votes2[i]
        ALP_vote_perc = in_flow.df$Vote_perc2[i]
        
        LNP_votes = in_flow.df$Votes1[i]
        LNP_vote_perc = in_flow.df$Vote_perc1[i]
      }
      
      candidate = in_flow.df$From_candidate[i]
      party = in_flow.df$From_party[i]
      party_file = replace_party_abbrev(party)
      calc_type = in_flow.df$Calc_type[i]
      
      this_out_line = sprintf("%d,%d,%s,%s,%s,%d,%.2f,%d,%.2f,%s",
                              year, i_year, party, seat, candidate, ALP_votes, ALP_vote_perc,
                              LNP_votes, LNP_vote_perc, calc_type)
      
      party_output_file = sprintf("party_flow_all_%s.csv", party_file)
      
      if (!(file.exists(party_output_file))) {
        write("Year,Election_i,Party,Seat,Candidate,ALP_votes,ALP_vote_perc,LNP_votes,LNP_vote_perc,Est_type",
              file=party_output_file)
      }
      
      write(this_out_line, file=party_output_file, append=TRUE)
    }
  }
}


##############################################################
# Part 2: Create summary files for each party, with mean and
# stdev of flow for each year.

party_files = list.files(path=".", pattern="party_flow_all.*\\.csv")
parties = gsub("party_flow_all_", "", party_files)
parties = gsub("\\.csv", "", parties)

num_elections = length(years)

base.df = data.frame(Year = years,
                     Election_i = seq_along(years),
                     stringsAsFactors = FALSE)

base.df$Count = 1
base.df$Perc_ALP = NA
base.df$StDev = NA

# The preference flows are extrapolated (and sometimes interpolated if there's too big
# a gap between observations) with a regressed mean, which is based on the following
# two sorta-eyeballed sorta-fudged parameters:
prior_sigma = 15
obs_sigma = 12

for (party in parties) {
  infile_flow = sprintf("party_flow_all_%s.csv", party)
  outfile_flow = sprintf("party_flow_summary_%s.csv", party)
  
  in_flow.df = read_csv(infile_flow)
  in_flow_by_year.df = group_by(in_flow.df, Year)
  
  summary_flow.df = summarise(in_flow_by_year.df,
                              Election_i = mean(Election_i),
                              Count = n(),
                              Mean_ALP = round(mean(ALP_vote_perc), 2),
                              SD_ALP = round(sd(ALP_vote_perc), 2))
 
  min_i = min(summary_flow.df$Election_i)
  max_i = max(summary_flow.df$Election_i)
  this_mean_ALP = mean(summary_flow.df$Mean_ALP)
  total_N = sum(summary_flow.df$Count)
  
  mu = c(50, this_mean_ALP)
  sigma = c(1/prior_sigma^2, total_N / obs_sigma^2)
  regressed_mean = round(sum(mu*sigma) / sum(sigma), 2)
  
  
  this.df = base.df
  this.df$Perc_ALP[summary_flow.df$Election_i] = summary_flow.df$Mean_ALP
  this.df$Count[summary_flow.df$Election_i] = summary_flow.df$Count
  this.df$Weight = this.df$Count
  this.df$StDev[summary_flow.df$Election_i] = summary_flow.df$SD_ALP
  
  this.df$Perc_ALP[which(this.df$Election_i < min_i)] = regressed_mean
  this.df$Perc_ALP[which(this.df$Election_i > max_i)] = regressed_mean
  
  num_obs = length(which(!(is.na(this.df$Perc_ALP))))
  
  if (num_obs < 7) {
    # Probably an early observation then a long gap before the next.
    this.df$Perc_ALP[which(is.na(this.df$Perc_ALP))] = regressed_mean
  }
  
  this_loess = loess(Perc_ALP ~ Election_i, data=this.df, weights = Weight, span=0.5)
  this.df$Perc_ALP_smoothed = round(predict(this_loess, data.frame(Election_i = seq_along(years))), 2)
  
  
  write.csv(this.df, file=outfile_flow, row.names=FALSE)
}


##############################################################
# Part 3: Estimate TPP's.  These are done in two parts: 
# estimating from preference flows where possible, then
# applying uniform swings in cases where there weren't
# candidates from both Labor and Coalition.

# First, estimating from flows.
party_flow_files = list.files(path=".", pattern="party_flow_summary.*\\.csv")
parties = gsub("\\.csv", "", party_flow_files)
parties = gsub("party_flow_summary_", "", parties)

outfile_ignored = "unestimated_tpp.csv"
write("Year,Seat,Parties", file=outfile_ignored)

for (i in seq_along(parties))  {
  infile_flow = party_flow_files[i]
  this.df = read_csv(infile_flow)
  this.df$Party = parties[i]
  
  if (i == 1) {
    party_flows.df = this.df
  } else {
    party_flows.df = rbind(party_flows.df, this.df)
  }
}

seat_states.df = read_csv("../misc/seat_states.csv")

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)

if (use_minimal_data == 1) {
  years = c(years, 1983, 1984, 1987, 1990, 1993, 1996, 1998, 2001, 2004, 2007, 2010, 2013)
}

for (year in years) {
  infile_tcp = sprintf("%s/%s_%d.csv", exact_folder, tcp_name(year), year)
  infile_prim = sprintf("%s/prim_%d.csv", exact_folder, year)
  infile_prefs = sprintf("%s/prefs_%d.csv", exact_folder, year)
  
  outfile_tcp = sprintf("tcp_%d.csv", year)
  outfile_missing = sprintf("missing_flow_%d.csv", year)
  
  write("Seat,Candidate,Party", file=outfile_missing)
  
  tcp.df = read_csv(infile_tcp)
  prim.df = read_csv(infile_prim)
  prefs.df = read_csv(infile_prefs)
  
  party_flows_year.df = filter(party_flows.df, Year == year)
  
  tcp.df$State = ""
  tcp.df$TPP_ALP = 0
  tcp.df$TPP_ALP_perc = 0
  tcp.df$TPP_LNP = 0
  tcp.df$TPP_LNP_perc = 0
  tcp.df$TPP_est_type = ""
  
  for (i in seq_along(tcp.df$Seat)) {
    seat = tcp.df$Seat[i]
    tcp.df$State[i] = seat_states.df$State[which(seat_states.df$Seat == seat)]
    
    if ((year == 1987) && (seat == "Blaxland")) {
      # Skip
    } else {
      seat_prim.df = filter(prim.df, Seat == seat)
      seat_prefs.df = filter(prefs.df, Seat == seat)
      
      tcp_parties = c(tcp.df$Party1[i], tcp.df$Party2[i])
      tpp_i = tpp_indices(tcp_parties)
      
      do_estimate = 0
      check_primaries = 0
      ignored = 0
      
      if (tcp.df$Est[i] == "Psephos") {
        # Don't have exact TCP.
        if (tpp_i[1] == 1) {
          # TPP contest
          do_estimate = 1
          
          if (tpp_i[2] == 1) {
            ALP_cand = tcp.df$Candidate1[i]
            LNP_cand = tcp.df$Candidate2[i]
          } else {
            ALP_cand = tcp.df$Candidate2[i]
            LNP_cand = tcp.df$Candidate1[i]
          }
        } else {
          # Non-classic.
          check_primaries = 1
        }
      } else {
        # Have exact TCP, but check if it's non-classic.
        if (tpp_i[1] == 0) {
          check_primaries = 1
        } else {
          # Classic division.
          if (tpp_i[2] == 1) {
            tcp.df$TPP_ALP[i] = tcp.df$Votes1[i]
            tcp.df$TPP_LNP[i] = tcp.df$Votes2[i]
            tcp.df$TPP_ALP_perc[i] = tcp.df$Vote_perc1[i]
            tcp.df$TPP_LNP_perc[i] = tcp.df$Vote_perc2[i]
          } else {
            tcp.df$TPP_ALP[i] = tcp.df$Votes2[i]
            tcp.df$TPP_LNP[i] = tcp.df$Votes1[i]
            tcp.df$TPP_ALP_perc[i] = tcp.df$Vote_perc2[i]
            tcp.df$TPP_LNP_perc[i] = tcp.df$Vote_perc1[i]
          }
          tcp.df$TPP_est_type[i] = "exact"
        }
      }
      
      if (check_primaries == 1) {
        any_tpp_list = any_tpp(seat_prim.df$Party)
        
        if (any_tpp_list$tpp_possible == 1) {
          # TPP estimate is possible.
          # Get the highest-polling candidates for ALP and LNP.
          
          do_estimate = 1
          
          ALP_i = any_tpp_list$ALP
          LNP_i = any_tpp_list$LNP
          
          ALP_prim.df = seat_prim.df[ALP_i, ]
          LNP_prim.df = seat_prim.df[LNP_i, ]
          
          temp_ALP_cand_i = which(ALP_prim.df$Votes == max(ALP_prim.df$Votes))
          temp_LNP_cand_i = which(LNP_prim.df$Votes == max(LNP_prim.df$Votes))
          
          ALP_cand = seat_prim.df$Candidate[ALP_i[temp_ALP_cand_i]]
          LNP_cand = seat_prim.df$Candidate[LNP_i[temp_LNP_cand_i]]
        } else {
          ignored = 1
          out_list_parties = paste(seat_prim.df$Party, collapse=";")
        }
      }
      
      if (do_estimate == 1) {
        # Here we gooooo!
        
        # In the original version of this code, the estimates were done
        # straight off the primaries.  Instead, check to see if there
        # was a partial preference distribution.  If there was, then
        # use the latest count containing both the LNP and ALP candidates.
        
        use_prim = TRUE
        
        if (length(seat_prefs.df$Seat) > 0) {
          max_count = max(seat_prefs.df$Count)
          for (count_i in max_count:2) {
            count.df = filter(seat_prefs.df, Count == count_i)
            ALP_i = which(count.df$To_candidate == ALP_cand)
            LNP_i = which(count.df$To_candidate == LNP_cand)
            
            if (length(ALP_i)*length(LNP_i) == 1) {
              use_prim = FALSE
              last_exact.df = transmute(count.df,
                                        Seat = Seat,
                                        Candidate = To_candidate,
                                        Party = To_party,
                                        Votes = New_votes,
                                        Vote_perc = New_vote_perc)
              break
            }
          }
        }
        
        if (use_prim) {
          last_exact.df = seat_prim.df
        }
        
        ALP_i = which(last_exact.df$Candidate == ALP_cand)
        LNP_i = which(last_exact.df$Candidate == LNP_cand)
        
        ALP_votes = last_exact.df$Votes[ALP_i]
        LNP_votes = last_exact.df$Votes[LNP_i]
        
        other_last_exact.df = last_exact.df[-c(ALP_i, LNP_i), ]
        
        for (j in seq_along(other_last_exact.df$Party)) {
          this_votes = other_last_exact.df$Votes[j]
          party = replace_party_abbrev(other_last_exact.df$Party[j])
          party = replace_party_for_flow(year, seat, other_last_exact.df$Candidate[j], party)
          
          party_i = which(party_flows_year.df$Party == party)
          
          if (length(party_i) == 0) {
            this_ALP_flow = 0.5
            
            this_out_line = sprintf("%s,%s,%s", seat, other_last_exact.df$Candidate[j], party)
            write(this_out_line, file=outfile_missing, append=TRUE)
          } else {
            if (length(party_i) != 1) {
              stop("Bad length_party_i")
            }
            
            this_ALP_flow = party_flows_year.df$Perc_ALP_smoothed[party_i] / 100
            
            if (party == "Ind") {
              this_ALP_flow = 0.5
            }
          }
          
          ALP_votes = ALP_votes + this_votes * this_ALP_flow
          LNP_votes = LNP_votes + this_votes * (1 - this_ALP_flow)
        }
        
        ALP_perc = ALP_votes / (ALP_votes + LNP_votes)
        LNP_perc = 1 - ALP_perc
        
        # Do some rounding here to ensure that the
        # rounded ALP_votes + LNP_votes = this_total_votes
        this_total_votes = round(ALP_votes + LNP_votes)
        LNP_votes = this_total_votes - round(ALP_votes)
        
        tcp.df$TPP_ALP[i] = round(ALP_votes)
        tcp.df$TPP_LNP[i] = round(LNP_votes)
        tcp.df$TPP_ALP_perc[i] = round(100*ALP_perc, 2)
        tcp.df$TPP_LNP_perc[i] = round(100*LNP_perc, 2)
        tcp.df$TPP_est_type[i] = "estimate"
      }
      
      if (ignored == 1) {
        write(sprintf("%s,%s,%s", year, seat, out_list_parties), file=outfile_ignored, append=TRUE)
      }
    }
  }
  
  # May be useful to inspect this in a spreadsheet:
  # write.csv(tcp.df, file=outfile_tcp, row.names=FALSE)
  
  tcp.df$Year = year
  i_year = which(years == year)
  tcp.df$Election_i = i_year
  
  if (year == years[1]) {
    all_tcp.df = tcp.df
  } else {
    all_tcp.df = rbind(all_tcp.df, tcp.df)
  }
}

# Kludge for 1943 Newcastle, which doesn't appear in the TCP tables because
# the order of exclusion of candidates couldn't be determined (also, there
# was no Coalition candidate):
write("1943,Newcastle,various", file=outfile_ignored, append=TRUE)

# Now the guessing of TPP's based on applying uniform swings.

# Add the unopposed divisions to the ignored file
infile_unopposed = "unopposed_divisions.csv"
unopposed.df = read_csv(infile_unopposed)

for (i in seq_along(unopposed.df$Seat)) {
  seat = unopposed.df$Seat[i]
  year = unopposed.df$Year[i]
  
  write(sprintf("%s,%s", year, seat), file=outfile_ignored, append=TRUE)
}

# And load all of the unestimated divisions into a new data frame:
unestimated.df = read_csv(outfile_ignored)
unestimated.df = unestimated.df[with(unestimated.df, order(Year)), ]
unestimated.df$TPP_ALP_perc = 0
unestimated.df$State = ""

for (i in seq_along(unestimated.df$Seat)) {
  seat = unestimated.df$Seat[i]
  unestimated.df$State[i] = seat_states.df$State[which(seat_states.df$Seat == seat)]
}

# Get TPP swings from the seats in common between subsequent elections.
outfile_tpp_guesses = "tpp_guesses.csv"

all_tpp.df = filter(all_tcp.df, TPP_ALP > 0)
all_tpp.df = all_tpp.df[c("Year", "State", "Seat", "TPP_ALP", "TPP_ALP_perc", "TPP_LNP", "TPP_LNP_perc", "TPP_est_type")]

# National swings:
# swings = get_swings(all_tpp.df, years, "TPP_ALP", "TPP_LNP")
# unestimated.df = apply_swings(all_tpp.df, unestimated.df, years, swings, "TPP_ALP_perc")

# State swings:
swings.df = get_state_swings(all_tpp.df, years, "TPP_ALP", "TPP_LNP")
unestimated.df = apply_state_swings(all_tpp.df, unestimated.df, years, swings.df, "TPP_ALP_perc")

# That gives us guessed TPP percentages for these seats.  We need also to estimate
# the formal votes.
unestimated.df$Formal_perc = 0

# Load all the turnout figures:
for (i in seq_along(years)) {
  year = years[i]
  infile_turnout = sprintf("%s/turnout_%d.csv", exact_folder, year)
  turnout.df = read_csv(infile_turnout)
  turnout.df$Year = year
  
  if (i == 1) {
    all_turnout.df = turnout.df
  } else {
    all_turnout.df = rbind(all_turnout.df, turnout.df)
  }
}

all_turnout.df$Formal = with(all_turnout.df, Voted - Informal)
all_turnout.df$Non_formal = with(all_turnout.df, Enrolled - Formal)
all_turnout.df$Formal_perc = with(all_turnout.df, 100 * Formal / (Formal + Non_formal))

# Compulsory voting was introduced in 1925, so do the votes estimate in two parts.
non_compulsory_years = years[which(years < 1925)]
compulsory_years = years[which(years >= 1925)]

non_compulsory_swings = get_swings(all_turnout.df, non_compulsory_years, "Formal", "Non_formal")
unestimated.df = apply_swings(all_turnout.df,
                              unestimated.df,
                              non_compulsory_years,
                              non_compulsory_swings,
                              "Formal_perc")


compulsory_swings = get_swings(all_turnout.df, compulsory_years, "Formal", "Non_formal")
unestimated.df = apply_swings(all_turnout.df,
                              unestimated.df,
                              compulsory_years,
                              compulsory_swings,
                              "Formal_perc")

# Fill in missing formal vote percentages (caused by, e.g. Warringah only having
# one election before compulsory voting, and the candidate being unopposed) with
# the overall formal vote percentage for that election.

missing_formal_i = which(is.na(unestimated.df$Formal_perc))
missing_formal_years = unique(unestimated.df$Year[missing_formal_i])

for (year in missing_formal_years) {
  turnout.df = filter(all_turnout.df, Year == year)
  turnout.df = filter(turnout.df, !(grepl("Territory", Seat) & (Year < 1966)))
  
  mean_formal_perc = 100 * (sum(turnout.df$Voted) - sum(turnout.df$Informal)) / sum(turnout.df$Enrolled)
  
  this_missing_i = intersect(which(unestimated.df$Year == year), missing_formal_i)
  unestimated.df$Formal_perc[this_missing_i] = mean_formal_perc
}

# Now we have formal votes as a percentage of enrolment.  Need to read off the 
# enrolment figure.
unestimated.df$Enrolled = 0

for (i in seq_along(unestimated.df$Seat)) {
  year = unestimated.df$Year[i]
  seat = unestimated.df$Seat[i]
  
  # Check first in the turnout data.
  this_turnout.df = filter(all_turnout.df, (Seat == seat) & (Year == year))
  if (length(this_turnout.df$Seat) > 0) {
    unestimated.df$Enrolled[i] = this_turnout.df$Enrolled[1]
  } else {
    # Look in the unopposed data.
    this_unopposed.df = filter(unopposed.df, (Seat == seat) & (Year == year))
    unestimated.df$Enrolled[i] = this_unopposed.df$Enrolled[1]
  }
}

unestimated.df = mutate(unestimated.df,
                        Formal_votes = round(Enrolled * Formal_perc / 100),
                        Formal_perc = round(Formal_perc, 2),
                        TPP_ALP = round(Formal_votes * TPP_ALP_perc / 100),
                        TPP_LNP = Formal_votes - TPP_ALP,
                        TPP_ALP_perc = round(TPP_ALP_perc, 2),
                        TPP_LNP_perc = 100 - TPP_ALP_perc)

unestimated.df = unestimated.df[c("Year", "State", "Seat", "Enrolled", "Formal_perc",
                                  "TPP_ALP", "TPP_ALP_perc", "TPP_LNP", "TPP_LNP_perc")]

write.csv(unestimated.df, file=outfile_tpp_guesses, row.names=FALSE)

# Put them all together.
unestimated.df = unestimated.df[c("Year", "State", "Seat", "TPP_ALP", "TPP_ALP_perc",
                                  "TPP_LNP", "TPP_LNP_perc")]

unestimated.df$TPP_est_type = "guess"

all_tpp.df = rbind(all_tpp.df, unestimated.df)

all_tpp.df = arrange(all_tpp.df, Year, Seat)

years = unique(all_tpp.df$Year)

# Delete the files made last time this script was run:
prev_tpp_files = list.files(path=".", pattern="tpp_[0-9]{4}\\.csv")
file.remove(prev_tpp_files)

for (year in years) {
  outfile = sprintf("tpp_%d.csv", year)
  
  year_tpp.df = filter(all_tpp.df, Year == year)
  write.csv(year_tpp.df, file=outfile, row.names=FALSE)
  
  if ((use_minimal_data == 1) && (year > 1980)) {
    # Scatter plot true v estimated TPP's.
    infile_tcp = sprintf("../exact_data/%s_%d.csv", tcp_name(year), year)
    
    exact_tcp.df = read_csv(infile_tcp)
    
    est_seats = year_tpp.df$Seat[which(year_tpp.df$TPP_est_type == "estimate")]
    exact_tcp.df = filter(exact_tcp.df, Seat %in% est_seats)
    
    year_tpp.df = filter(year_tpp.df, Seat %in% est_seats)
    year_tpp.df$Exact_TPP_ALP_perc = 0
    year_tpp.df$Exact_TPP_LNP_perc = 0
    
    for (i in seq_along(year_tpp.df$Seat)) {
      seat = year_tpp.df$Seat[i]
      match_i = which(exact_tcp.df$Seat == seat)
      
      if (length(match_i) != 1) {
        stop("Couldn't match TPP")
      }
      
      tcp_parties = c(exact_tcp.df$Party1[match_i], exact_tcp.df$Party2[match_i])
      tpp_i = tpp_indices(tcp_parties)
      
      if (tpp_i[1] == 0) {
        print(sprintf("Non-classic %d %s", year, seat))
      } else {
        if (tpp_i[2] == 1) {
          year_tpp.df$Exact_TPP_ALP_perc[i] = exact_tcp.df$Vote_perc1[match_i]
          year_tpp.df$Exact_TPP_LNP_perc[i] = exact_tcp.df$Vote_perc2[match_i]
        } else {
          year_tpp.df$Exact_TPP_ALP_perc[i] = exact_tcp.df$Vote_perc2[match_i]
          year_tpp.df$Exact_TPP_LNP_perc[i] = exact_tcp.df$Vote_perc1[match_i]
        }
      }
    }
    
    outfile_tpp = sprintf("true_v_est_tpp_%d.csv", year)
    write.csv(year_tpp.df, file=outfile_tpp, row.names=FALSE)
    
    plot_img = ggplot() + geom_point(data=year_tpp.df,
                                     aes(x=TPP_ALP_perc, y=Exact_TPP_ALP_perc))
    
    outfile_img = sprintf("true_v_est_tpp_%d.png", year)
    
    png_print(plot_img, outfile_img, width=700, height=700)
    
    rho = cor(year_tpp.df$TPP_ALP_perc, year_tpp.df$Exact_TPP_ALP_perc)
    err = year_tpp.df$TPP_ALP_perc - year_tpp.df$Exact_TPP_ALP_perc
    
    print(sprintf("%d: r = %.4f, mean err = %.2f, mean abs err = %.2f",
                  year, rho, mean(err), mean(abs(err))))
  }
}

all_tpp.df = filter(all_tpp.df, !(grepl("Territory", Seat) & (Year < 1966)))

election_tpp.df = summarise(group_by(all_tpp.df, Year),
                            ALP = sum(TPP_ALP),
                            LNP = sum(TPP_LNP),
                            TPP_ALP_perc = round(100 * ALP / (ALP + LNP), 2))

options(warn=1)

