get_party_groups = function() {
  # Parties grouped together for preference-flow averaging purposes.
  
  party_groups = list()
  party_groups[[1]] = list(party="CP", abbrevs=c("NPA", "FSA", "VFU", "QFU", "CP", "LCP",
                                                 "NCP75", "PPU", "UCP22", "UCP34", "ACP34", "CPP"))
  party_groups[[2]] = list(party="Prog", abbrevs=c("WP", "Prog"))
  party_groups[[3]] = list(party="DLP", abbrevs=c("DLP", "QLP", "ALPAC"))
  party_groups[[4]] = list(party="LL", abbrevs=c("LL31", "LL40", "LL43"))
  party_groups[[5]] = list(party="Soc", abbrevs=c("SPA72", "SWP", "ISLP", "SLL"))
  party_groups[[6]] = list(party="Nat", abbrevs=c("Nat", "Lib22"))
  party_groups[[7]] = list(party="UAP31", abbrevs=c("UAP31", "EC"))
  party_groups[[8]] = list(party="Ind UAP31", abbrevs=c("Ind UAP31", "Ind EC", "Ind Lib46"))
  party_groups[[9]] = list(party="Ind ALP", abbrevs=c("Ind ALP", "PLP22"))
  
  return(party_groups)
}


replace_party_abbrev = function(input_party) {
  # Much like replace_party_abbrevs() in merge_exact.R, but here
  # it only works for a single input_party
  
  party_groups = get_party_groups()
  
  num_replace_abbrevs = length(party_groups)
  
  for (i in 1:num_replace_abbrevs) {
    if (input_party %in% party_groups[[i]]$abbrevs) {
      input_party = party_groups[[i]]$party
      return(input_party)
    }
  }
  
  return(input_party)
}

replace_party_for_flow = function(year, seat, cand, party) {
  # From around the 1960's, independent candidates who were aligned with a
  # party stopped being referred to as, e.g., "Ind Lab".  This function 
  # makes some such replacements for the purposes of allocating their
  # preferences.
  #
  # There is no Ind Lib46 preferences data to work with, so Independent
  # Liberals are coded here as Ind Nat.
  
  out_party = party
  
  if ((year == 1966) && (seat == "Batman") && (cand == "Sam Benson")) {
    out_party = "Ind ALP"
  }
  
  if ((year == 1966) && (seat == "Richmond") && (cand == "Keith Compton")) {
    out_party = "Ind ALP"
  }
  
  if ((year == 1966) && (seat == "Warringah") && (cand == "Keith Chambers")) {
    out_party = "Ind Nat"
  }
  
  if ((year == 1969) && (seat == "Murray") && (cand == "William Hunter")) {
    out_party = "Ind Nat"
  }
  
  if ((year == 1969) && (seat == "Sydney") && (cand == "Nick Origlass")) {
    out_party = "CPA"
  }
  
  if ((year == 1969) && (seat == "Warringah") && (cand == "Edward St John")) {
    out_party = "Ind Nat"
  }
  
  if ((year == 1972) && (seat == "Australian Capital Territory") && (cand == "Arthur Burns")) {
    out_party = "Ind ALP"
  }
  
  return(out_party)
}


tpp_groups = function() {
  tpp_list = list(lnp = c("Lib46", "NPA", "CLP", "LNP", "FSA", "VFU", "QFU", "CP", "EC", "LCP",
                          "Nat", "UAP31", "NCP75", "PPU", "Lib22", "UCP22", "UCP34", "ACP34", "CPP"),
                  alp = c("ALP", "LL31", "LL40", "LL43", "SL"))
  
  return(tpp_list)
}

tpp_indices = function(parties) {
  # Input a two-element vector of parties.
  # Returns a three-element vector:
  #   First element is 1 if it is a TPP contest
  #   Second element is the index of the ALP candidate
  #   Third element is the index of the LNP candidate
  
  tpp_parties = tpp_groups()
  
  coalition_parties = tpp_parties$lnp
  labor_parties = tpp_parties$alp
  
  classic = 0
  
  # In the following, ALP_i and LNP_i will usually point to Labor and
  # Coalition candidates respectively.  But they might point to an
  # independent (or whatever); the idea is to make the order consistent
  # in the output CSV files.
  
  if ((parties[1] %in% labor_parties) || (parties[2] %in% labor_parties)){
    ALP_i = which(parties %in% labor_parties)[1]
    LNP_i = 3 - ALP_i
    
    if (parties[LNP_i] %in% coalition_parties) {
      classic = 1
    }
  } else {
    LNP_i = which(parties %in% coalition_parties)
    if (length(LNP_i) == 0) {
      # Neither LNP nor ALP are in last two.
      LNP_i = 1
    }
    
    LNP_i = LNP_i[1]
    ALP_i = 3 - LNP_i
  }
  
  return(c(classic, ALP_i, LNP_i))
}

any_tpp = function(parties) {
  # Input is a vector of parties
  # Returns a list containing the indices of Labor candidates, and the indices
  # of Coalition candidates
  
  tpp_parties = tpp_groups()
  
  coalition_parties = tpp_parties$lnp
  labor_parties = tpp_parties$alp
  
  LNP_i = which(parties %in% coalition_parties)
  ALP_i = which(parties %in% labor_parties)
  
  if (length(LNP_i)*length(ALP_i) > 0) {
    tpp_possible = 1
  } else {
    tpp_possible = 0
  }
  
  output_list = list(tpp_possible = tpp_possible,
                     ALP = ALP_i,
                     LNP = LNP_i)
  
  return(output_list)
}

tcp_name = function(year) {
  if (year < 1993) {
    return("tcp")
  } else {
    return("tpp")
  }
}

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)
}


apply_swings = function(known.df, unknown.df, years, swings, varname) {
  # Input data frames assumed to have fields "Seat", "Year", varname.
  # Outputs unknown.df with the guessed varname's (expect this to be
  # either TPP_ALP_perc or Formal_vote_perc).
  
  all_years = years
  
  # When calling this function to estimate vote totals, it's important
  # to separate out years of compulsory and non-compulsory voting.
  # Temporarily remove rows from the data frames from the years
  # being ignored:
  known_all_years.df = known.df
  unknown_all_years.df = unknown.df
  
  known_i = which(known.df$Year %in% years)
  unknown_i = which(unknown.df$Year %in% years)
  
  known.df = known.df[known_i, ]
  unknown.df = unknown.df[unknown_i, ]
  
  seats = unique(unknown.df$Seat)
  
  for (seat in seats) {
    years = all_years
    
    seat_known.df = filter(known.df, Seat == seat)
    seat_unknown.df = filter(unknown.df, Seat == seat)
    seat_unknown_i = which(unknown.df$Seat == seat)
    
    # Truncate the years, otherwise if a seat was abolished in (say)
    # 1975, the script will try to back-estimate 1977 and 1980 from
    # non-existent data.
    #
    # I haven't quite worked out why similar problems don't exist
    # on the other side, i.e., for electorates created later than 1919.
    max_year = max(c(seat_known.df$Year, seat_unknown.df$Year))
    
    years = all_years[which(all_years <= max_year)]
    
    var_fwd = 0*years
    var_back = 0*years
    var_est_count = 0*years
    
    year_i = which(years %in% seat_known.df$Year)
    new_year_i = which(years %in% seat_unknown.df$Year)
    
    
    if (length(seat_known.df[[varname]]) == 0) {
      print(sprintf("%s %s", varname, seat))
    }
    
    var_fwd[year_i] = seat_known.df[[varname]]
    var_back[year_i] = seat_known.df[[varname]]
    var_est_count[year_i] = 2
    
    # Apply swings forwards.  Start in 1922; have to skip i = 1.
    for (i in new_year_i) {
      if (i > 1) {
        if ((var_fwd[i-1] > 0) && (var_fwd[i] == 0)) {
          # Careful with the indices:
          #   swings[1] is the swing in 1922,
          #   years[1] is 1919,
          #   ALP_TPP_fwd[1] is the ALP TPP in 1919.
          
          var_fwd[i] = var_fwd[i-1] + swings[i-1]
          var_est_count[i] = var_est_count[i] + 1
        }
      }
    }
    
    # Apply swings backwards.
    for (i in rev(new_year_i)) {
      if (i < length(years)) {
        if ((var_back[i+1] > 0) && (var_back[i] == 0)) {
          # Careful with the indices:
          #   swings[1] is the swing in 1922 from 1919,
          #   years[1] is 1919,
          #   ALP_TPP_back[1] is the ALP TPP in 1919.
          
          var_back[i] = var_back[i+1] - swings[i]
          var_est_count[i] = var_est_count[i] + 1
        }
      }
    }
    
    var_avg = (var_fwd + var_back) / (var_est_count)
    
    unknown.df[[varname]][seat_unknown_i] = var_avg[new_year_i]
  }
  
  unknown_all_years.df[[varname]][unknown_i] = unknown.df[[varname]]
  
  return(unknown_all_years.df)
}

apply_state_swings = function(known.df, unknown.df, years, swings.df, varname) {
  # As for apply_swings, except the swings.df input is a data frame with
  # state swings.
  
  all_years = years
  
  # When calling this function to estimate vote totals, it's important
  # to separate out years of compulsory and non-compulsory voting.
  # Temporarily remove rows from the data frames from the years
  # being ignored:
  known_all_years.df = known.df
  unknown_all_years.df = unknown.df
  
  known_i = which(known.df$Year %in% years)
  unknown_i = which(unknown.df$Year %in% years)
  
  known.df = known.df[known_i, ]
  unknown.df = unknown.df[unknown_i, ]
  
  seats = unique(unknown.df$Seat)
  
  for (seat in seats) {
    years = all_years
    
    seat_known.df = filter(known.df, Seat == seat)
    seat_unknown.df = filter(unknown.df, Seat == seat)
    seat_unknown_i = which(unknown.df$Seat == seat)
    
    state = seat_unknown.df$State[1]
    if (grepl("(ACT|NT)", state)) {
      state = "Aust"
    }
    
    
    # Truncate the years, otherwise if a seat was abolished in (say)
    # 1975, the script will try to back-estimate 1977 and 1980 from
    # non-existent data.
    #
    # I haven't quite worked out why similar problems don't exist
    # on the other side, i.e., for electorates created later than 1919.
    max_year = max(c(seat_known.df$Year, seat_unknown.df$Year))
    
    years = all_years[which(all_years <= max_year)]
    
    var_fwd = 0*years
    var_back = 0*years
    var_est_count = 0*years
    
    year_i = which(years %in% seat_known.df$Year)
    new_year_i = which(years %in% seat_unknown.df$Year)
    
    
    if (length(seat_known.df[[varname]]) == 0) {
      print(sprintf("%s %s", varname, seat))
    }
    
    var_fwd[year_i] = seat_known.df[[varname]]
    var_back[year_i] = seat_known.df[[varname]]
    var_est_count[year_i] = 2
    
    # Apply swings forwards.  Start in 1922; have to skip i = 1.
    for (i in new_year_i) {
      if (i > 1) {
        if ((var_fwd[i-1] > 0) && (var_fwd[i] == 0)) {
          # Careful with the indices:
          #   swings[1] is the swing in 1922,
          #   years[1] is 1919,
          #   ALP_TPP_fwd[1] is the ALP TPP in 1919.
          
          var_fwd[i] = var_fwd[i-1] + swings.df[[state]][i-1]
          var_est_count[i] = var_est_count[i] + 1
        }
      }
    }
    
    # Apply swings backwards.
    for (i in rev(new_year_i)) {
      if (i < length(years)) {
        if ((var_back[i+1] > 0) && (var_back[i] == 0)) {
          # Careful with the indices:
          #   swings[1] is the swing in 1922 from 1919,
          #   years[1] is 1919,
          #   ALP_TPP_back[1] is the ALP TPP in 1919.
          
          var_back[i] = var_back[i+1] - swings.df[[state]][i]
          var_est_count[i] = var_est_count[i] + 1
        }
      }
    }
    
    var_avg = (var_fwd + var_back) / (var_est_count)
    
    unknown.df[[varname]][seat_unknown_i] = var_avg[new_year_i]
  }
  
  unknown_all_years.df[[varname]][unknown_i] = unknown.df[[varname]]
  
  return(unknown_all_years.df)
}


get_swings = function(known.df, years, numer_var, denom_var) {
  # Variable of interest is numer_var / (numer_var + denom_var),
  # e.g., TPP_ALP / (TPP_ALP + TPP_LNP)
  #
  # known.df assumed to have fields Seat, Year, numer_var, denom_var
  
  # Exclude the territories from the swing calculations
  # when they didn't really count:
  known.df = filter(known.df, !(grepl("Territory", Seat) & (Year < 1966)))
  
  swings = 0 * seq_along(years)[-1]
  
  for (i in seq_along(years)[-1]) {
    year1 = years[i-1]
    year2 = years[i]
    
    year1.df = filter(known.df, Year == year1)
    year2.df = filter(known.df, Year == year2)
    
    common_seats = intersect(year1.df$Seat, year2.df$Seat)
    
    year1.df = filter(year1.df, Seat %in% common_seats)
    year2.df = filter(year2.df, Seat %in% common_seats)
    
    var1 = 100 * sum(year1.df[[numer_var]]) / (sum(year1.df[[numer_var]]) + sum(year1.df[[denom_var]]))
    var2 = 100 * sum(year2.df[[numer_var]]) / (sum(year2.df[[numer_var]]) + sum(year2.df[[denom_var]]))
    
    swings[i-1] = var2 - var1
  }
  
  return(swings)
}

get_state_swings = function(known.df, years, numer_var, denom_var) {
  # get_swings() but returns a data frame with the swings by state.
  
  known.df = filter(known.df, !(grepl("Territory", Seat) & (Year < 1966)))
  
  swings = -101 * seq_along(years)[-1]
  swings.df = data.frame(Aust = swings)
  states = c("NSW", "Vic", "Qld", "WA", "SA", "Tas")
  
  for (state in states) {
    swings.df[[state]] = swings
  }
  
  states = c("Aust", states)
  
  for (i in seq_along(years)[-1]) {
    year1 = years[i-1]
    year2 = years[i]
    
    for (state in states) {
      if (state == "Aust") {
        year1.df = filter(known.df, Year == year1)
        year2.df = filter(known.df, Year == year2)
      } else {
        year1.df = filter(known.df, (Year == year1) & (State == state))
        year2.df = filter(known.df, (Year == year2) & (State == state))
      }
      
      common_seats = intersect(year1.df$Seat, year2.df$Seat)
      
      year1.df = filter(year1.df, Seat %in% common_seats)
      year2.df = filter(year2.df, Seat %in% common_seats)
      
      var1 = 100 * sum(year1.df[[numer_var]]) / (sum(year1.df[[numer_var]]) + sum(year1.df[[denom_var]]))
      var2 = 100 * sum(year2.df[[numer_var]]) / (sum(year2.df[[numer_var]]) + sum(year2.df[[denom_var]]))
      
      swings.df[[state]][i-1] = var2 - var1
    }
  }
  
  return(swings.df)
}