OpenTally/src/main.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/>.
 */

use opentally::stv;
use opentally::election::{Candidate, CandidateState, CountCard, CountState, CountStateOrRef, Election, StageResult};
use opentally::numbers::{NativeFloat64, Number, Rational};

use clap::{AppSettings, Clap};
use git_version::git_version;

use std::fs::File;
use std::io::{self, BufRead};
use std::ops;

const VERSION: &str = git_version!(args=["--always", "--dirty=-dev"], fallback="unknown");

/// Open-source election vote counting
#[derive(Clap)]
#[clap(name="OpenTally", version=VERSION)]
struct Opts {
	#[clap(subcommand)]
	command: Command,
}

#[derive(Clap)]
enum Command {
	STV(STV),
}

/// Count a single transferable vote (STV) election
#[derive(Clap)]
#[clap(setting=AppSettings::DeriveDisplayOrder)]
struct STV {
	// ----------------
	// -- File input --
	
	/// Path to the BLT file to be counted
	filename: String,
	
	// ----------------------
	// -- Numbers settings --
	
	/// Numbers mode
	#[clap(help_heading=Some("NUMBERS"), short, long, possible_values=&["rational", "float64"], default_value="rational", value_name="mode")]
	numbers: String,
	
	// -----------------------
	// -- Rounding settings --
	
	/// Round transfer values to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_tvs: Option<usize>,
	
	/// Round ballot weights to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_weights: Option<usize>,
	
	/// Round votes to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_votes: Option<usize>,
	
	/// Round quota to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_quota: Option<usize>,
	
	// ------------------
	// -- STV variants --
	
	/// Method of surplus transfers
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["wig", "uig", "eg", "meek"], default_value="wig", value_name="method")]
	surplus: String,
	
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["by_size", "by_order"], default_value="by_size", value_name="order")]
	surplus_order: String,
	
	/// Examine only transferable papers during surplus distributions
	#[clap(help_heading=Some("STV VARIANTS"), long)]
	transferable_only: bool,
	
	/// Method of exclusions
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["single_stage", "by_value", "parcels_by_order"], default_value="single_stage", value_name="method")]
	exclusion: String,
	
	// ----------------------
	// -- Display settings --
	
	/// Hide excluded candidates from results report
	#[clap(help_heading=Some("DISPLAY"), long)]
	hide_excluded: bool,
	
	/// Sort candidates by votes in results report
	#[clap(help_heading=Some("DISPLAY"), long)]
	sort_votes: bool,
	
	/// Print votes to specified decimal places in results report
	#[clap(help_heading=Some("DISPLAY"), long, default_value="2", value_name="dps")]
	pp_decimals: usize,
}

fn main() {
	// Read arguments
	let opts: Opts = Opts::parse();
	let Command::STV(cmd_opts) = opts.command;
	
	// Read BLT file
	let file = File::open(&cmd_opts.filename).expect("IO Error");
	let lines = io::BufReader::new(file).lines();
	
	// Create and count election according to --numbers
	if cmd_opts.numbers == "rational" {
		let election: Election<Rational> = Election::from_blt(lines.map(|r| r.expect("IO Error").to_string()).into_iter());
		count_election(election, cmd_opts);
	} else if cmd_opts.numbers == "float64" {
		let election: Election<NativeFloat64> = Election::from_blt(lines.map(|r| r.expect("IO Error").to_string()).into_iter());
		count_election(election, cmd_opts);
	}
}

fn count_election<N: Number>(election: Election<N>, cmd_opts: STV)
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	// Copy applicable options
	let stv_opts = stv::STVOptions::new(
		cmd_opts.round_tvs,
		cmd_opts.round_weights,
		cmd_opts.round_votes,
		cmd_opts.round_quota,
		&cmd_opts.surplus,
		&cmd_opts.surplus_order,
		cmd_opts.transferable_only,
		&cmd_opts.exclusion,
		cmd_opts.pp_decimals,
	);
	
	// Initialise count state
	let mut state = CountState::new(&election);
	
	// Distribute first preferences
	stv::count_init(&mut state, &stv_opts);
	let mut stage_num = 1;
	make_and_print_result(stage_num, &state, &cmd_opts);
	
	loop {
		let is_done = stv::count_one_stage(&mut state, &stv_opts);
		if is_done {
			break;
		}
		stage_num += 1;
		make_and_print_result(stage_num, &state, &cmd_opts);
	}
	
	println!("Count complete. The winning candidates are, in order of election:");
	
	let mut winners = Vec::new();
	for (candidate, count_card) in state.candidates.iter() {
		if count_card.state == CandidateState::ELECTED {
			winners.push((candidate, count_card.order_elected));
		}
	}
	winners.sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
	
	for (i, (winner, _)) in winners.into_iter().enumerate() {
		println!("{}. {}", i + 1, winner.name);
	}
}

fn print_candidates<'a, N: 'a + Number, I: Iterator<Item=(&'a Candidate, &'a CountCard<'a, N>)>>(candidates: I, cmd_opts: &STV) {
	for (candidate, count_card) in candidates {
		if count_card.state == CandidateState::ELECTED {
			println!("- {}: {:.dps$} ({:.dps$}) - ELECTED {}", candidate.name, count_card.votes, count_card.transfers, count_card.order_elected, dps=cmd_opts.pp_decimals);
		} else if count_card.state == CandidateState::EXCLUDED {
			// If --hide-excluded, hide unless nonzero votes or nonzero transfers
			if !cmd_opts.hide_excluded || !count_card.votes.is_zero() || !count_card.transfers.is_zero() {
				println!("- {}: {:.dps$} ({:.dps$}) - Excluded {}", candidate.name, count_card.votes, count_card.transfers, -count_card.order_elected, dps=cmd_opts.pp_decimals);
			}
		} else {
			println!("- {}: {:.dps$} ({:.dps$})", candidate.name, count_card.votes, count_card.transfers, dps=cmd_opts.pp_decimals);
		}
	}
}

fn print_stage<N: Number>(stage_num: usize, result: &StageResult<N>, cmd_opts: &STV) {
	// Print stage details
	match result.kind {
		None => { println!("{}. {}", stage_num, result.title); }
		Some(kind) => { println!("{}. {} {}", stage_num, kind, result.title); }
	};
	println!("{}", result.logs.join(" "));
	
	let state = result.state.as_ref();
	
	// Print candidates
	if cmd_opts.sort_votes {
		// Sort by votes if requested
		let mut candidates: Vec<(&Candidate, &CountCard<N>)> = state.candidates.iter()
			.map(|(c, cc)| (*c, cc)).collect();
		// First sort by order of election (as a tie-breaker, if votes are equal)
		candidates.sort_unstable_by(|a, b| b.1.order_elected.partial_cmp(&a.1.order_elected).unwrap());
		// Then sort by votes
		candidates.sort_by(|a, b| a.1.votes.partial_cmp(&b.1.votes).unwrap());
		print_candidates(candidates.into_iter().rev(), cmd_opts);
	} else {
		let candidates = state.election.candidates.iter()
			.map(|c| (c, state.candidates.get(c).unwrap()));
		print_candidates(candidates, cmd_opts);
	}
	
	// Print summary rows
	println!("Exhausted: {:.dps$} ({:.dps$})", state.exhausted.votes, state.exhausted.transfers, dps=cmd_opts.pp_decimals);
	println!("Loss by fraction: {:.dps$} ({:.dps$})", state.loss_fraction.votes, state.loss_fraction.transfers, dps=cmd_opts.pp_decimals);
	
	let mut total_vote = state.candidates.values().fold(N::zero(), |acc, cc| { acc + &cc.votes });
	total_vote += &state.exhausted.votes;
	total_vote += &state.loss_fraction.votes;
	println!("Total votes: {:.dps$}", total_vote, dps=cmd_opts.pp_decimals);
	
	println!("Quota: {:.dps$}", state.quota, dps=cmd_opts.pp_decimals);
	
	println!("");
}

fn make_and_print_result<N: Number>(stage_num: usize, state: &CountState<N>, cmd_opts: &STV) {
	let result = StageResult {
		kind: state.kind,
		title: &state.title,
		logs: state.logger.render(),
		state: CountStateOrRef::from(&state),
	};
	print_stage(stage_num, &result, &cmd_opts);
}