Implement eSTV-style CSV report

src/cli/stv.rs

@@ -16,13 +16,14 @@
  */
 
 use crate::constraints::Constraints;
-use crate::election::{CandidateState, CountState, Election};
+use crate::election::{CandidateState, CountState, Election, StageKind};
 use crate::numbers::{Fixed, GuardedFixed, NativeFloat64, Number, Rational};
 use crate::parser::{bin, blt};
 use crate::stv::{self, STVOptions};
 use crate::ties;
 
 use clap::{AppSettings, Clap};
+use itertools::Itertools;
 
 use std::cmp::max;
 use std::fs::File;
@@ -174,19 +175,22 @@ pub struct SubcmdOptions {
 	#[clap(help_heading=Some("CONSTRAINTS"), long, possible_values=&["guard_doom"], default_value="guard_doom")]
 	constraint_mode: String,
 	
-	// ----------------------
+	// ---------------------
-	// -- Display settings --
+	// -- Output settings --
+	
+	#[clap(help_heading=Some("OUTPUT"), short, long, possible_values=&["text", "csv"], default_value="text")]
+	output: String,
 	
 	/// Hide excluded candidates from results report
-	#[clap(help_heading=Some("DISPLAY"), long)]
+	#[clap(help_heading=Some("OUTPUT"), long)]
 	hide_excluded: bool,
 	
 	/// Sort candidates by votes in results report
-	#[clap(help_heading=Some("DISPLAY"), long)]
+	#[clap(help_heading=Some("OUTPUT"), 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")]
+	#[clap(help_heading=Some("OUTPUT"), long, default_value="2", value_name="dps")]
 	pp_decimals: usize,
 }
 
@@ -244,7 +248,7 @@ fn maybe_load_constraints<N: Number>(election: &mut Election<N>, constraints: &O
 	}
 }
 
-fn count_election<N: Number>(mut election: Election<N>, cmd_opts: SubcmdOptions) -> Result<(), i32>
+fn count_election<N: Number>(election: Election<N>, cmd_opts: SubcmdOptions) -> Result<(), i32>
 where
 	for<'r> &'r N: ops::Add<&'r N, Output=N>,
 	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
@@ -293,9 +297,27 @@ where
 		}
 	}
 	
+	match cmd_opts.output.as_str() {
+		"text" => { return count_election_text(election, &cmd_opts.filename, opts); }
+		"csv" => { return count_election_csv(election, opts); }
+		_ => unreachable!()
+	}
+}
+
+// ---------------
+// CLI text report
+
+fn count_election_text<N: Number>(mut election: Election<N>, filename: &str, opts: STVOptions) -> Result<(), i32>
+where
+	for<'r> &'r N: ops::Add<&'r N, Output=N>,
+	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
+	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
+	for<'r> &'r N: ops::Div<&'r N, Output=N>,
+	for<'r> &'r N: ops::Neg<Output=N>
+{
 	// Describe count
 	let total_ballots = election.ballots.iter().fold(N::zero(), |acc, b| { acc + &b.orig_value });
-	print!("Count computed by OpenTally (revision {}). Read {:.0} ballots from \"{}\" for election \"{}\". There are {} candidates for {} vacancies. ", crate::VERSION, total_ballots, cmd_opts.filename, election.name, election.candidates.len(), election.seats);
+	print!("Count computed by OpenTally (revision {}). Read {:.0} ballots from \"{}\" for election \"{}\". There are {} candidates for {} vacancies. ", crate::VERSION, total_ballots, filename, election.name, election.candidates.len(), election.seats);
 	let opts_str = opts.describe::<N>();
 	if opts_str.len() > 0 {
 		println!("Counting using options \"{}\".", opts_str);
@@ -359,7 +381,7 @@ where
 	return Ok(());
 }
 
-fn print_stage<N: Number>(stage_num: usize, state: &CountState<N>, opts: &STVOptions) {
+fn print_stage<N: Number>(stage_num: u32, state: &CountState<N>, opts: &STVOptions) {
 	// Print stage details
 	println!("{}. {}", stage_num, state.title);
 	println!("{}", state.logger.render().join(" "));
@@ -372,3 +394,175 @@ fn print_stage<N: Number>(stage_num: usize, state: &CountState<N>, opts: &STVOpt
 	
 	println!("");
 }
+
+// ---------------------
+// eSTV-style CSV report
+
+fn count_election_csv<N: Number>(mut election: Election<N>, opts: STVOptions) -> Result<(), i32>
+where
+	for<'r> &'r N: ops::Add<&'r N, Output=N>,
+	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
+	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
+	for<'r> &'r N: ops::Div<&'r N, Output=N>,
+	for<'r> &'r N: ops::Neg<Output=N>
+{
+	// Header rows
+	let total_ballots = election.ballots.iter().fold(N::zero(), |acc, b| { acc + &b.orig_value });
+	
+	// eSTV does not consistently quote records, so we won't use a CSV library here
+	println!(r#""Election for","{}""#, election.name);
+	println!(r#""Date","  /  /    ""#);
+	println!(r#""Number to be elected",{}"#, election.seats);
+	
+	stv::preprocess_election(&mut election, &opts);
+	
+	// Initialise count state
+	let mut state = CountState::new(&election);
+	
+	let mut stage_results = vec![Vec::new(); election.candidates.len() + 5];
+	
+	// -----------
+	// First stage
+	
+	// Distribute first preferences
+	match stv::count_init(&mut state, &opts) {
+		Ok(_) => {}
+		Err(err) => {
+			println!("Error: {}", err.describe());
+			return Err(1);
+		}
+	}
+	
+	// Subtract this from progressive NTs
+	// FIXME: May fail to round correctly with minivoters
+	let invalid_votes = state.exhausted.votes.clone();
+	let valid_votes = total_ballots - &invalid_votes;
+	
+	// Stage number row
+	stage_results[0].push(String::new());
+	stage_results[0].push(String::new());
+	
+	// Stage kind row
+	stage_results[1].push(String::new());
+	stage_results[1].push(String::from(r#""First""#));
+	
+	// Stage title row
+	stage_results[2].push(String::from(r#""Candidates""#));
+	stage_results[2].push(String::from(r#""Preferences""#));
+	
+	for (i, candidate) in election.candidates.iter().enumerate() {
+		let count_card = &state.candidates[candidate];
+		stage_results[3 + i].push(format!(r#""{}""#, candidate.name));
+		stage_results[3 + i].push(format!(r#"{:.0}"#, count_card.votes)); // FIXME: May fail to round correctly with minivoters
+	}
+	
+	stage_results[3 + election.candidates.len()].push(String::from(r#""Non-transferable""#));
+	stage_results[3 + election.candidates.len()].push(String::new()); // FIXME: May fail to round correctly with minivoters
+	
+	stage_results[4 + election.candidates.len()].push(String::from(r#""Totals""#));
+	stage_results[4 + election.candidates.len()].push(format!(r#"{:.0}"#, valid_votes));
+	
+	// -----------------
+	// Subsequent stages
+	
+	let mut stage_num: u32 = 1;
+	loop {
+		match stv::count_one_stage(&mut state, &opts) {
+			Ok(is_done) => {
+				if is_done {
+					break;
+				}
+			}
+			Err(err) => {
+				println!("Error: {}", err.describe());
+				return Err(1);
+			}
+		}
+		
+		stage_num += 1;
+		
+		// Stage number row
+		stage_results[0].push(String::from(r#""Stage""#));
+		stage_results[0].push(format!(r#"{}"#, stage_num));
+		
+		// Stage kind row
+		stage_results[1].push(format!(r#""{}""#, state.title.kind_as_string()));
+		stage_results[1].push(String::new());
+		
+		// Stage title row
+		match &state.title {
+			StageKind::FirstPreferences => unreachable!(),
+			StageKind::SurplusOf(candidate) => {
+				stage_results[2].push(format!(r#""{}""#, candidate.name));
+			}
+			StageKind::ExclusionOf(candidates) => {
+				stage_results[2].push(format!(r#""{}""#, candidates.iter().map(|c| &c.name).join("+")));
+			}
+			StageKind::SurplusesDistributed => todo!(),
+			StageKind::BulkElection => todo!(),
+		}
+		stage_results[2].push(String::from(r#""#));
+		
+		for (i, candidate) in election.candidates.iter().enumerate() {
+			let count_card = &state.candidates[candidate];
+			
+			if count_card.transfers.is_zero() {
+				stage_results[3 + i].push(String::new());
+			} else if count_card.transfers > N::zero() {
+				stage_results[3 + i].push(format!(r#"+{:.dps$}"#, count_card.transfers, dps=opts.pp_decimals));
+			} else {
+				stage_results[3 + i].push(format!(r#"{:.dps$}"#, count_card.transfers, dps=opts.pp_decimals));
+			}
+			
+			if count_card.votes.is_zero() {
+				stage_results[3 + i].push(String::from(r#""-""#));
+			} else {
+				stage_results[3 + i].push(format!(r#"{:.dps$}"#, count_card.votes, dps=opts.pp_decimals));
+			}
+		}
+		
+		// Nontransferable
+		let nt_transfers = state.exhausted.transfers.clone() + &state.loss_fraction.transfers;
+		if nt_transfers.is_zero() {
+			stage_results[3 + election.candidates.len()].push(String::new());
+		} else if nt_transfers > N::zero() {
+			stage_results[3 + election.candidates.len()].push(format!(r#"+{:.dps$}"#, nt_transfers, dps=opts.pp_decimals));
+		} else {
+			stage_results[3 + election.candidates.len()].push(format!(r#"{:.dps$}"#, nt_transfers, dps=opts.pp_decimals));
+		}
+		stage_results[3 + election.candidates.len()].push(format!(r#"{:.dps$}"#, &state.exhausted.votes + &state.loss_fraction.votes, dps=opts.pp_decimals));
+		
+		// Totals
+		stage_results[4 + election.candidates.len()].push(String::new());
+		stage_results[4 + election.candidates.len()].push(format!(r#"{:.dps$}"#, valid_votes, dps=opts.pp_decimals));
+	}
+	
+	// ----------------
+	// Candidate states
+	
+	stage_results[3 + election.candidates.len()].push(String::new()); // Nontransferable row
+	
+	for (i, candidate) in election.candidates.iter().enumerate() {
+		let count_card = &state.candidates[candidate];
+		if count_card.state == CandidateState::Elected {
+			stage_results[3 + i].push(String::from(r#""Elected""#));
+		} else {
+			stage_results[3 + i].push(String::new());
+		}
+	}
+	
+	// --------------------
+	// Output stages to CSV
+	
+	println!(r#""Valid votes",{:.0}"#, valid_votes);
+	println!(r#""Invalid votes",{:.0}"#, invalid_votes);
+	println!(r#""Quota",{:.dps$}"#, state.quota.as_ref().unwrap(), dps=opts.pp_decimals);
+	println!(r#""OpenTally","{}""#, crate::VERSION);
+	println!(r#""Election rules","{}""#, opts.describe::<N>());
+	
+	for row in stage_results {
+		println!("{}", row.join(","));
+	}
+	
+	return Ok(());
+}