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/>.
 */

mod election;
mod numbers;
mod stv;

use crate::election::{CandidateState, CountState, CountStateOrRef, Election, StageResult};
use crate::numbers::{Number, NumType};

use clap::Clap;

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

#[derive(Clap)]
struct Opts {
	#[clap(subcommand)]
	command: Command,
}

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

/// Count a single transferable vote (STV) election
#[derive(Clap)]
struct STV {
	/// Path to the BLT file to be counted
	filename: String,
	/// Print votes to specified decimal places in results report
	#[clap(long, default_value="2")]
	pp_decimals: usize,
}

fn print_stage<N: Number>(stage_num: usize, result: &StageResult<N>, cmd_opts: &STV) {
	println!("{}. {}", stage_num, result.title);
	println!("{}", result.logs.join(" "));
	
	// Sort candidates
	//let mut candidates: Vec<(&&Candidate, &CountCard<N>)> = result.state.candidates.iter().collect();
	//candidates.sort_unstable_by(|a, b| a.1.votes.partial_cmp(&b.1.votes).unwrap());
	let candidates = result.state.as_ref().election.candidates.iter().map(|c| (c, result.state.as_ref().candidates.get(c).unwrap()));
	
	// Print candidates
	//for (candidate, count_card) in candidates.into_iter().rev() {
	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 {
			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);
		}
	}
	
	println!("Quota: {:.dps$}", result.state.as_ref().quota, dps=cmd_opts.pp_decimals);
	
	println!("");
}

fn main() {
	let should_clone_state = false;
	
	// 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();
	let election: Election<NumType> = Election::from_blt(lines);
	
	// Initialise count state
	let mut state = CountState::new(&election);
	
	// Distribute first preferences
	stv::distribute_first_preferences(&mut state);
	stv::calculate_quota(&mut state);
	stv::elect_meeting_quota(&mut state);
	
	// Display
	// TODO: Add logs during count
	let result = StageResult {
		title: "First preferences",
		logs: vec!["First preferences distributed."],
		state: if should_clone_state { CountStateOrRef::State(state.clone()) } else { CountStateOrRef::Ref(&state) },
	};
	print_stage(1, &result, &cmd_opts);
	
	let mut stage_num = 1;
	
	loop {
		state.step_all();
		stage_num += 1;
		
		// Finish count
		if stv::finished_before_stage(&state) {
			break;
		}
		
		// Continue exclusions
		if stv::continue_exclusion(&mut state) {
			stv::elect_meeting_quota(&mut state);
			let result = StageResult {
				title: "Exclusion",
				logs: vec!["Continuing exclusion."],
				state: if should_clone_state { CountStateOrRef::State(state.clone()) } else { CountStateOrRef::Ref(&state) },
			};
			print_stage(stage_num, &result, &cmd_opts);
			continue;
		}
		
		// Distribute surpluses
		if stv::distribute_surpluses(&mut state) {
			stv::elect_meeting_quota(&mut state);
			let result = StageResult {
				title: "Surplus",
				logs: vec!["Surplus distributed."],
				state: if should_clone_state { CountStateOrRef::State(state.clone()) } else { CountStateOrRef::Ref(&state) },
			};
			print_stage(stage_num, &result, &cmd_opts);
			continue;
		}
		
		// Attempt bulk election
		if stv::bulk_elect(&mut state) {
			stv::elect_meeting_quota(&mut state);
			let result = StageResult {
				title: "Bulk election",
				logs: vec!["Bulk election."],
				state: if should_clone_state { CountStateOrRef::State(state.clone()) } else { CountStateOrRef::Ref(&state) },
			};
			print_stage(stage_num, &result, &cmd_opts);
			continue;
		}
		
		// Exclude lowest hopeful
		if stv::exclude_hopefuls(&mut state) {
			stv::elect_meeting_quota(&mut state);
			let result = StageResult {
				title: "Exclusion",
				logs: vec!["Candidate excluded."],
				state: if should_clone_state { CountStateOrRef::State(state.clone()) } else { CountStateOrRef::Ref(&state) },
			};
			print_stage(stage_num, &result, &cmd_opts);
			continue;
		}
		
		todo!();
	}
	
	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);
	}
}