path: root/src/stv/gregory/transfers.rs

/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021 Lee Yingtong Li (RunasSudo)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#[cfg(not(target_arch = "wasm32"))]
use prettytable::{Cell, Row, Table};
#[cfg(target_arch = "wasm32")]
use super::prettytable_html::{Cell, Row, Table};

use crate::election::{Candidate, CountState};
use crate::numbers::Number;
use crate::stv::{STVOptions, SumSurplusTransfersMode};

use std::collections::HashMap;

/// Table describing vote transfers during a surplus distribution or exclusion
pub struct TransferTable<'e, N: Number> {
	/// Continuing candidates
	pub hopefuls: Vec<&'e Candidate>,
	
	/// Columns in the table
	pub columns: Vec<TransferTableColumn<'e, N>>,
	
	/// Total column
	pub total: TransferTableColumn<'e, N>,
	
	/// Size of surplus, or `None` if an exclusion
	pub surplus: Option<N>,
	/// Surplus fraction, or `None` if votes not reweighted/an exclusion (for display/optimisation only)
	pub surpfrac: Option<N>,
	/// Numerator of surplus fraction, or `None` if votes not reweighted/an exclusion
	pub surpfrac_numer: Option<N>,
	/// Denominator of surplus fraction, or `None`
	pub surpfrac_denom: Option<N>,
}

impl<'e, N: Number> TransferTable<'e, N> {
	/// Return a new [TransferTable] for an exclusion
	pub fn new_exclusion(hopefuls: Vec<&'e Candidate>) -> Self {
		TransferTable {
			hopefuls,
			columns: Vec::new(),
			total: TransferTableColumn {
				value_fraction: N::new(),
				cells: HashMap::new(),
				exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
				total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			},
			surplus: None,
			surpfrac: None,
			surpfrac_numer: None,
			surpfrac_denom: None,
		}
	}
	
	/// Return a new [TransferTable] for a surplus distribution
	pub fn new_surplus(hopefuls: Vec<&'e Candidate>, surplus: N, surpfrac: Option<N>, surpfrac_numer: Option<N>, surpfrac_denom: Option<N>) -> Self {
		TransferTable {
			hopefuls,
			columns: Vec::new(),
			total: TransferTableColumn {
				value_fraction: N::new(),
				cells: HashMap::new(),
				exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
				total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			},
			surplus: Some(surplus),
			surpfrac,
			surpfrac_numer,
			surpfrac_denom,
		}
	}
	
	/// Record the specified transfer
	pub fn add_transfers(&mut self, value_fraction: &N, candidate: &'e Candidate, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction {
				col.add_transfers(candidate, ballots);
				return;
			}
		}
		
		let mut col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			cells: HashMap::new(),
			exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
		};
		col.add_transfers(candidate, ballots);
		self.columns.push(col);
	}
	
	/// Record the specified exhaustion
	pub fn add_exhausted(&mut self, value_fraction: &N, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction {
				col.exhausted.ballots += ballots;
				return;
			}
		}
		
		let col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			cells: HashMap::new(),
			exhausted: TransferTableCell { ballots: ballots.clone(), votes_in: N::new(), votes_out: N::new() },
			total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
		};
		self.columns.push(col);
	}
	
	/// Calculate the votes to be transferred according to this table
	pub fn calculate(&mut self, opts: &STVOptions) {
		// Use weighted rules if exclusion or WIGM
		let is_weighted = self.surplus.is_none() || opts.surplus.is_weighted();
		
		// Iterate through columns
		for column in self.columns.iter_mut() {
			let mut new_value_fraction = N::new();
			if self.surplus.is_some() && opts.sum_surplus_transfers == SumSurplusTransfersMode::PerBallot {
				if is_weighted {
					new_value_fraction = column.value_fraction.clone();
					// If surplus, multiply by surplus fraction
					if let Some(n) = &self.surpfrac_numer {
						new_value_fraction *= n;
					}
				} else {
					if let Some(n) = &self.surpfrac_numer {
						new_value_fraction = n.clone();
					} else {
						// Transferred at original value
						new_value_fraction = column.value_fraction.clone();
					}
				}
				
				if let Some(n) = &self.surpfrac_denom {
					new_value_fraction /= n;
				}
				
				// Round if required
				if let Some(dps) = opts.round_values {
					new_value_fraction.floor_mut(dps);
				}
			}
			
			// Candidate votes
			for (candidate, cell) in column.cells.iter_mut() {
				column.total.ballots += &cell.ballots;
				self.total.add_transfers(*candidate, &cell.ballots);
				self.total.total.ballots += &cell.ballots;
				
				let votes_in = cell.ballots.clone() * &column.value_fraction;
				cell.votes_in += &votes_in;
				column.total.votes_in += &votes_in;
				self.total.cells.get_mut(*candidate).unwrap().votes_in += &votes_in;
				self.total.total.votes_in += votes_in;
				
				if self.surplus.is_some() && opts.sum_surplus_transfers == SumSurplusTransfersMode::PerBallot {
					let votes_out = cell.ballots.clone() * &new_value_fraction;
					cell.votes_out += &votes_out;
					column.total.votes_out += &votes_out;
					self.total.cells.get_mut(*candidate).unwrap().votes_out += &votes_out;
					self.total.total.votes_out += votes_out;
				}
			}
			
			// Exhausted votes
			column.total.ballots += &column.exhausted.ballots;
			self.total.exhausted.ballots += &column.exhausted.ballots;
			self.total.total.ballots += &column.exhausted.ballots;
			
			let votes_in = column.exhausted.ballots.clone() * &column.value_fraction;
			column.exhausted.votes_in += &votes_in;
			column.total.votes_in += &votes_in;
			self.total.exhausted.votes_in += &votes_in;
			self.total.total.votes_in += votes_in;
			
			if self.surplus.is_some() && opts.sum_surplus_transfers == SumSurplusTransfersMode::PerBallot {
				if !opts.transferable_only {
					let votes_out = column.exhausted.ballots.clone() * &new_value_fraction;
					column.exhausted.votes_out += &votes_out;
					column.total.votes_out += &votes_out;
					self.total.exhausted.votes_out += &votes_out;
					self.total.total.votes_out += votes_out;
				}
			}
		}
		
		// Need to calculate votes_out?
		if self.surplus.is_none() || opts.sum_surplus_transfers == SumSurplusTransfersMode::ByValue {
			for (_candidate, cell) in self.total.cells.iter_mut() {
				let mut votes_out;
				if is_weighted {
					votes_out = cell.votes_in.clone();
				} else {
					votes_out = cell.ballots.clone();
				}
				
				// If surplus, multiply by surplus fraction
				if let Some(n) = &self.surpfrac_numer {
					votes_out *= n;
				}
				if let Some(n) = &self.surpfrac_denom {
					votes_out /= n;
				}
				
				cell.votes_out = votes_out; // Rounded later
			}
			
			if self.surplus.is_none() || !opts.transferable_only {
				let mut votes_out;
				if is_weighted {
					votes_out = self.total.exhausted.votes_in.clone();
				} else {
					votes_out = self.total.exhausted.ballots.clone();
				}
				
				// If surplus, multiply by surplus fraction
				if let Some(n) = &self.surpfrac_numer {
					votes_out *= n;
				}
				if let Some(n) = &self.surpfrac_denom {
					votes_out /= n;
				}
				
				self.total.exhausted.votes_out = votes_out; // Rounded later
			}
		}
		
		// Round if required
		if let Some(dps) = opts.round_votes {
			for (_candidate, cell) in self.total.cells.iter_mut() {
				cell.votes_out.floor_mut(dps);
			}
			
			self.total.exhausted.votes_out.floor_mut(dps);
		}
	}
	
	/// Apply the transfers described in the table to the count sheet
	///
	/// Credit continuing candidates and exhausted pile with the appropriate number of ballot papers and votes.
	pub fn apply_to(&self, state: &mut CountState<N>, opts: &STVOptions) -> N {
		let mut checksum = N::new();
		
		// Credit transferred votes
		for (candidate, count_card) in state.candidates.iter_mut() {
			if let Some(cell) = self.total.cells.get(candidate) {
				count_card.transfer(&cell.votes_out);
				count_card.ballot_transfers += &cell.ballots;
				checksum += &cell.votes_out;
			}
		}
		
		// Credit exhausted votes
		// If exclusion or not --transferable-only
		if self.surplus.is_none() || !opts.transferable_only {
			// Standard rules
			state.exhausted.transfer(&self.total.exhausted.votes_out);
			state.exhausted.ballot_transfers += &self.total.exhausted.ballots;
			checksum += &self.total.exhausted.votes_out;
		} else {
			// Credit only nontransferable difference
			if self.surpfrac_numer.is_none() {
				// TODO: Is there a purer way of calculating this?
				let difference = self.surplus.as_ref().unwrap().clone() - &checksum;
				state.exhausted.transfer(&difference);
				checksum += difference;
				
				for column in self.columns.iter() {
					state.exhausted.ballot_transfers += &column.exhausted.ballots;
				}
			} else {
				// No ballots exhaust
			}
		}
		
		return checksum;
	}
	
	/// Render table as [Table]
	fn render(&self, opts: &STVOptions) -> Table {
		let mut table = Table::new();
		set_table_format(&mut table);
		
		let show_transfers_per_ballot = self.surpfrac.is_some() && opts.sum_surplus_transfers == SumSurplusTransfersMode::PerBallot;
		
		let num_cols;
		if show_transfers_per_ballot {
			num_cols = self.columns.len() * 3 + 4;
		} else {
			if self.surpfrac.is_none() {
				num_cols = self.columns.len() * 2 + 3;
			} else {
				num_cols = self.columns.len() * 2 + 4;
			}
		}
		
		// ----------
		// Header row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Preference"));
		for column in self.columns.iter() {
			row.push(Cell::new(&format!("Ballots @ {:.dps$}", column.value_fraction, dps=opts.pp_decimals)).style_spec("cH2"));
			
			if show_transfers_per_ballot {
				row.push(Cell::new(&format!("× {:.dps$}", self.surpfrac.as_ref().unwrap(), dps=opts.pp_decimals)).style_spec("r"));
			}
		}
		row.push(Cell::new("Total").style_spec("cH2"));
		if self.surpfrac.is_some() {
			row.push(Cell::new(&format!("× {:.dps$}", self.surpfrac.as_ref().unwrap(), dps=opts.pp_decimals)).style_spec("r"));
		}
		table.set_titles(Row::new(row));
		
		// --------------
		// Candidate rows
		
		for candidate in self.hopefuls.iter() {
			let mut row = Vec::with_capacity(num_cols);
			row.push(Cell::new(&candidate.name));
			for column in self.columns.iter() {
				if let Some(cell) = column.cells.get(candidate) {
					row.push(Cell::new(&format!("{:.0}", cell.ballots)).style_spec("r"));
					row.push(Cell::new(&format!("{:.dps$}", cell.votes_in, dps=opts.pp_decimals)).style_spec("r"));
					if show_transfers_per_ballot {
						row.push(Cell::new(&format!("{:.dps$}", cell.votes_out, dps=opts.pp_decimals)).style_spec("r"));
					}
				} else {
					row.push(Cell::new(""));
					row.push(Cell::new(""));
					if show_transfers_per_ballot {
						row.push(Cell::new(""));
					}
				}
			}
			
			// Totals
			if let Some(cell) = self.total.cells.get(candidate) {
				row.push(Cell::new(&format!("{:.0}", cell.ballots)).style_spec("r"));
				row.push(Cell::new(&format!("{:.dps$}", cell.votes_in, dps=opts.pp_decimals)).style_spec("r"));
				if self.surpfrac.is_some() {
					row.push(Cell::new(&format!("{:.dps$}", cell.votes_out, dps=opts.pp_decimals)).style_spec("r"));
				}
			} else {
				row.push(Cell::new(""));
				row.push(Cell::new(""));
				if self.surpfrac.is_some() {
					row.push(Cell::new(""));
				}
			}
			
			table.add_row(Row::new(row));
		}
		
		// -------------
		// Exhausted row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Exhausted"));
		for column in self.columns.iter() {
			if !column.exhausted.ballots.is_zero() {
				row.push(Cell::new(&format!("{:.0}", column.exhausted.ballots)).style_spec("r"));
				row.push(Cell::new(&format!("{:.dps$}", column.exhausted.votes_in, dps=opts.pp_decimals)).style_spec("r"));
				if show_transfers_per_ballot {
					if column.exhausted.votes_out.is_zero() {
						row.push(Cell::new("-").style_spec("c"));
					} else {
						row.push(Cell::new(&format!("{:.dps$}", column.exhausted.votes_out, dps=opts.pp_decimals)).style_spec("r"));
					}
				}
			} else {
				row.push(Cell::new(""));
				row.push(Cell::new(""));
				if show_transfers_per_ballot {
					row.push(Cell::new(""));
				}
			}
		}
		
		// Totals
		if !self.total.exhausted.ballots.is_zero() {
			row.push(Cell::new(&format!("{:.0}", self.total.exhausted.ballots)).style_spec("r"));
			row.push(Cell::new(&format!("{:.dps$}", self.total.exhausted.votes_in, dps=opts.pp_decimals)).style_spec("r"));
			if self.surpfrac.is_some() {
				if self.total.exhausted.votes_out.is_zero() {
					row.push(Cell::new("-").style_spec("c"));
				} else {
					row.push(Cell::new(&format!("{:.dps$}", self.total.exhausted.votes_out, dps=opts.pp_decimals)).style_spec("r"));
				}
			}
		} else {
			row.push(Cell::new(""));
			row.push(Cell::new(""));
			if self.surpfrac.is_some() {
				row.push(Cell::new(""));
			}
		}
		
		table.add_row(Row::new(row));
		
		// ----------
		// Totals row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Total"));
		
		for column in self.columns.iter() {
			row.push(Cell::new(&format!("{:.0}", column.total.ballots)).style_spec("r"));
			row.push(Cell::new(&format!("{:.dps$}", column.total.votes_in, dps=opts.pp_decimals)).style_spec("r"));
			if show_transfers_per_ballot {
				row.push(Cell::new(&format!("{:.dps$}", column.total.votes_out, dps=opts.pp_decimals)).style_spec("r"));
			}
		}
		
		// Grand total cell
		
		let mut gt_ballots = N::new();
		let mut gt_votes_in = N::new();
		let mut gt_votes_out = N::new();
		
		for candidate in self.hopefuls.iter() {
			if let Some(cell) = self.total.cells.get(candidate) {
				gt_ballots += &cell.ballots;
				gt_votes_in += &cell.votes_in;
				gt_votes_out += &cell.votes_out;
			}
		}
		gt_ballots += &self.total.exhausted.ballots;
		gt_votes_in += &self.total.exhausted.votes_in;
		gt_votes_out += &self.total.exhausted.votes_out;
		
		row.push(Cell::new(&format!("{:.0}", gt_ballots)).style_spec("r"));
		row.push(Cell::new(&format!("{:.dps$}", gt_votes_in, dps=opts.pp_decimals)).style_spec("r"));
		if self.surpfrac.is_some() {
			row.push(Cell::new(&format!("{:.dps$}", gt_votes_out, dps=opts.pp_decimals)).style_spec("r"));
		}
		
		table.add_row(Row::new(row));
		
		return table;
	}
	
	/// Render table as plain text
	#[cfg(not(target_arch = "wasm32"))]
	pub fn render_text(&self, opts: &STVOptions) -> String {
		return self.render(opts).to_string();
	}
	
	/// Render table as HTML
	#[cfg(target_arch = "wasm32")]
	pub fn render_text(&self, opts: &STVOptions) -> String {
		return self.render(opts).to_html();
	}
	
	// Render table as HTML
	//pub fn render_html(&self, opts: &STVOptions) -> String {
	//	return self.render(opts).to_string();
	//}
}

/// Column in a [TransferTable]
pub struct TransferTableColumn<'e, N: Number> {
	/// Value fraction of ballots counted in this column
	pub value_fraction: N,
	
	/// Cells in this column
	pub cells: HashMap<&'e Candidate, TransferTableCell<N>>,
	
	/// Exhausted cell
	pub exhausted: TransferTableCell<N>,
	
	/// Totals cell
	pub total: TransferTableCell<N>,
}

impl<'e, N: Number> TransferTableColumn<'e, N> {
	/// Record the specified transfer
	pub fn add_transfers(&mut self, candidate: &'e Candidate, ballots: &N) {
		if let Some(cell) = self.cells.get_mut(candidate) {
			cell.ballots += ballots;
		} else {
			let cell = TransferTableCell {
				ballots: ballots.clone(),
				votes_in: N::new(),
				votes_out: N::new(),
			};
			self.cells.insert(candidate, cell);
		}
	}
}

/// Cell in a [TransferTable], representing transfers to one candidate at a particular value
pub struct TransferTableCell<N: Number> {
	/// Ballots expressing a next preference for the continuing candidate
	pub ballots: N,
	/// Value of votes when received by the transferring candidate
	pub votes_in: N,
	/// Votes transferred to the continuing candidate
	pub votes_out: N,
}

#[cfg(not(target_arch = "wasm32"))]
fn set_table_format(table: &mut Table) {
	table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
}

#[cfg(target_arch = "wasm32")]
fn set_table_format(_table: &mut Table) {
	// No op
}