diff --git a/src/ballot.rs b/src/ballot.rs index 7e15082..06ee1a1 100644 --- a/src/ballot.rs +++ b/src/ballot.rs @@ -1,8 +1,10 @@ +use std::rc::Rc; + use itertools::Itertools; use crate::header::Header; #[derive(Copy,Clone,Debug,PartialEq)] -pub enum BallotType { +enum BallotType { Party, Candidate, } @@ -10,21 +12,18 @@ pub enum BallotType { #[derive(Debug)] pub struct Ballot { - pub state: String, - pub division: String, - pub collection_point: String, - pub collection_point_id: usize, - pub batch_id: usize, - pub paper_id: usize, - pub votes: Vec, - pub ty: BallotType, + header: Rc
, + collection_point_id: usize, + batch_id: usize, + paper_id: usize, + votes: Vec, } impl Ballot { - pub fn parse(row: quick_csv::Row, header: &Header) -> Result> { + pub fn parse(row: quick_csv::Row, header: Rc
) -> Result> { let mut cols = row.columns()?; - let (state, division, collection_point, collection_point_id, batch_id, paper_id) = cols.next_tuple().ok_or("Missing columns")?; + let (_, _, _, collection_point_id, batch_id, paper_id) = cols.next_tuple().ok_or("Missing columns")?; let cols = cols.collect_vec(); let mut votes: Vec>; let ty: BallotType; @@ -52,39 +51,34 @@ impl Ballot { } } + let votes = match ty { + BallotType::Candidate => votes.iter().filter_map(|&v| v).collect(), + BallotType::Party => votes.iter().filter_map(|&v| v).map(|id| header.parties[id].member_ids.iter().copied()).flatten().collect(), + }; + Ok(Ballot { - state: state.to_owned(), - division: division.to_owned(), - collection_point: collection_point.to_owned(), + header, collection_point_id: collection_point_id.parse()?, batch_id: batch_id.parse()?, paper_id: paper_id.parse()?, - votes: votes.iter().copied().filter_map(|v| v).collect(), - ty, + votes, }) } - - pub fn to_candidate_ballot(self, header: &Header) -> Self { - match self.ty { - BallotType::Candidate => self, - BallotType::Party => { - let mut votes = Vec::::new(); - - for party_id in self.votes { - votes.extend(header.parties[party_id].member_ids.iter().copied()); - } - - Ballot { - state: self.state, - division: self.division, - collection_point: self.collection_point, - collection_point_id: self.collection_point_id, - batch_id: self.batch_id, - paper_id: self.paper_id, - ty: BallotType::Candidate, - votes, - } - }, + pub fn get_primary_index(&self, enabled: &[bool]) -> Option { + for &id in self.votes.iter() { + if enabled[id] { + return Some(id); + } + } + return None; + } + pub fn count_all_with_decaying_weights(&self, scores: &mut [f64], enabled: &[bool]) { + let mut m = 0.5; + for &id in self.votes.iter() { + if enabled[id] { + scores[id] += m; + m *= 0.5; + } } } } diff --git a/src/candidate.rs b/src/candidate.rs index d34b0de..1688daa 100644 --- a/src/candidate.rs +++ b/src/candidate.rs @@ -1,3 +1,5 @@ +use crate::header::Header; + #[derive(Debug)] pub struct Candidate { @@ -5,3 +7,28 @@ pub struct Candidate { pub party_id: Option, } +#[derive(Copy,Clone)] +pub struct CandidateName<'a> { + candidate_name: &'a str, + party_name: Option<&'a str>, +} + +impl Candidate { + pub fn get_name<'a>(&'a self, header: &'a Header) -> CandidateName<'a> { + CandidateName { + candidate_name: &self.name, + party_name: self.party_id.map_or(None, |id| Some(&header.parties[id].name)), + } + } +} + +impl<'a> std::fmt::Display for CandidateName<'a> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.candidate_name)?; + if let Some(party_name) = self.party_name { + write!(f, ": {party_name}")?; + } + Ok(()) + } +} + diff --git a/src/counter.rs b/src/counter.rs index 458f5ff..5a6bfdc 100644 --- a/src/counter.rs +++ b/src/counter.rs @@ -1,9 +1,11 @@ -use crate::{ballot::{Ballot, BallotType}, header::Header}; +use std::rc::Rc; + +use crate::{ballot::Ballot, header::Header}; #[derive(Debug)] pub struct Counter { - pub header: Header, + pub header: Rc
, pub ballots: Vec, } @@ -13,7 +15,7 @@ impl Counter { let mut ballots = Vec::new(); for row in csv { - ballots.push(Ballot::parse(row?, &header)?.to_candidate_ballot(&header)); + ballots.push(Ballot::parse(row?, header.clone())?); } Ok(Counter { @@ -22,40 +24,27 @@ impl Counter { }) } - pub fn count_primaries(&self, ty: BallotType, running: &[bool]) -> Vec where T: num::Num + Copy + std::ops::AddAssign { - let mut scores = vec![T::zero(); running.len()]; + pub fn count_primaries(&self, enabled: &[bool]) -> Vec { + assert!(enabled.len() == self.header.candidates.len()); + + let mut scores = vec![0.0; enabled.len()]; for ballot in self.ballots.iter() { - if ballot.ty != ty { - continue; - } - for &id in ballot.votes.iter() { - if !running[id] { - continue; - } - scores[id] += T::one(); - break; + if let Some(index) = ballot.get_primary_index(enabled) { + scores[index] += 1.0; } } scores } - pub fn count_all_with_powers_of_2(&self, ty: BallotType, running: &[bool]) -> Vec { - let mut scores = vec![0.0; running.len()]; + pub fn count_all_with_decaying_weights(&self, enabled: &[bool]) -> Vec { + assert!(enabled.len() == self.header.candidates.len()); + + let mut scores = vec![0.0; enabled.len()]; for ballot in self.ballots.iter() { - if ballot.ty != ty { - continue; - } - let mut m = 0.5; - for &id in ballot.votes.iter() { - if !running[id] { - continue; - } - scores[id] += m; - m *= 0.5; - } + ballot.count_all_with_decaying_weights(&mut scores, enabled); } scores diff --git a/src/csv.rs b/src/csv.rs new file mode 100644 index 0000000..1c21d06 --- /dev/null +++ b/src/csv.rs @@ -0,0 +1,56 @@ +use crate::util::{EscapeWriter, EscapeWriterOpts}; +use std::io::Write; + + +pub struct CsvWriter where T: std::io::Write { + out: T, +} + +pub struct RowWriter<'a,T> where T: std::io::Write { + csv: &'a mut CsvWriter, + first: bool, +} + +impl CsvWriter where T: std::io::Write { + pub fn new(out: T) -> Self { + Self { out } + } + pub fn row(&mut self) -> RowWriter<'_,T> { + RowWriter { csv: self, first: true } + } +} + +impl<'a,T> RowWriter<'a,T> where T: std::io::Write { + pub fn skip(&mut self, cells: usize) -> std::io::Result<()> { + for _ in 0..cells { + if self.first { + self.first = false; + } + else { + write!(self.csv.out, ",")?; + } + } + Ok(()) + } + pub fn write(&mut self, value: F) -> std::io::Result<()> + where F: std::fmt::Display + { + self.skip(1)?; + write!(self.csv.out, "\"")?; + let mut writer = EscapeWriter::new(&mut self.csv.out, EscapeWriterOpts { + escape: b'\\', + chars: b"\"", + }); + write!(writer, "{value}")?; + write!(self.csv.out, "\"")?; + + Ok(()) + } +} + +impl<'a,T> Drop for RowWriter<'a,T> where T: std::io::Write { + fn drop(&mut self) { + _ = writeln!(self.csv.out); + } +} + diff --git a/src/header.rs b/src/header.rs index ec8896f..3980f1b 100644 --- a/src/header.rs +++ b/src/header.rs @@ -1,8 +1,8 @@ -use std::collections::HashMap; +use std::{collections::HashMap, rc::Rc}; use itertools::Itertools; -use crate::{candidate::Candidate, party::Party}; +use crate::{candidate::{Candidate, CandidateName}, party::Party}; #[derive(Debug)] @@ -12,7 +12,7 @@ pub struct Header { } impl Header { - pub fn parse(row: quick_csv::Row) -> Result> { + pub fn parse(row: quick_csv::Row) -> Result, Box> { let mut header = Header { parties: Vec::new(), candidates: Vec::new(), @@ -41,7 +41,13 @@ impl Header { } } - Ok(header) + Ok(header.into()) + } + pub fn get_candidate_name(&self, index: usize) -> CandidateName { + self.candidates[index].get_name(self) + } + pub fn iter_candidate_name(&self) -> impl Iterator { + self.candidates.iter().map(|v| v.get_name(self)) } } diff --git a/src/main.rs b/src/main.rs index 3532ac2..865de99 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,11 @@ use std::env; - -use ballot::BallotType; use counter::Counter; +use csv::CsvWriter; use itertools::Itertools; +use util::{percent::PercentDiff, Percent}; +pub mod csv; +pub mod util; pub mod candidate; pub mod ballot; pub mod header; @@ -21,94 +23,90 @@ fn main() { }; let csv = quick_csv::Csv::from_file(csv_path).unwrap().flexible(true); let counter = Counter::new(csv).unwrap(); + let header = counter.header.clone(); - eprintln!("Header: {:#?}", counter.header); - eprintln!("Parties: {}", counter.header.parties.len()); - eprintln!("Candidates: {}", counter.header.candidates.len()); + eprintln!("Header: {:#?}", header); + eprintln!("Parties: {}", header.parties.len()); + eprintln!("Candidates: {}", header.candidates.len()); let total = counter.ballots.len() as u32; - let mut total_atl = 0u32; - let mut total_btl = 0u32; + let mut running = vec![true; header.candidates.len()]; + let mut csv_out = CsvWriter::new(std::io::stdout().lock()); - for ballot in counter.ballots.iter() { - *match ballot.ty { - BallotType::Party => &mut total_atl, - BallotType::Candidate => &mut total_btl, - } += 1; + { + let mut row = csv_out.row(); + row.skip(4).unwrap(); + for name in header.iter_candidate_name() { + row.write(name).unwrap(); + row.skip(1).unwrap(); + } } - - eprintln!("Above the line: {total_atl} {}%", f64::from(total_atl) / f64::from(total) * 100.0); - eprintln!("Below the line: {total_btl} {}%", f64::from(total_btl) / f64::from(total) * 100.0); - eprintln!("Total: {total}"); - - let ty = BallotType::Candidate; - let names = match ty { - BallotType::Party => counter.header.parties.iter().map(|v| v.name.clone()).collect_vec(), - BallotType::Candidate => counter.header.candidates.iter().map(|v| match v.party_id { - Some(party_id) => format!("{}: {}", v.name, &counter.header.parties[party_id].name), - None => v.name.to_owned(), - }).collect_vec(), - }; - let mut running = vec![true; names.len()]; - - print!(",,,"); - for name in names.iter() { - print!(",{name:?},"); - } - println!(); eprintln!("Running election rounds:"); - let to_percent = |score: f64| score / f64::from(total) * 100.0; let mut scores_last: Option> = None; let mut id_worst_last: Option = None; for index in 0..(running.len() - winner_count + 1) { - let scores = counter.count_all_with_powers_of_2(ty, &running); - let scores_ordered = scores.iter().copied().enumerate().filter(|&(i,_)| running[i]).sorted_by(|a, b| f64::total_cmp(&b.1, &a.1)).collect_vec(); + let scores = counter.count_all_with_decaying_weights(&running); + let scores_ordered = scores.iter().copied().enumerate().filter(|&(id,_)| running[id]).sorted_by(|a, b| f64::total_cmp(&b.1, &a.1)).collect_vec(); + let (id_worst, score_worst) = match scores_ordered.last() { Some(&v) => v, None => break, }; + csv_out.row(); + let mut row = csv_out.row(); + if let Some(scores_last) = scores_last { let id_worst_last = id_worst_last.unwrap(); let score_worst_last = scores_last[id_worst_last]; - print!(",{:?},{score_worst_last},{:.3}%", names[id_worst_last], to_percent(score_worst_last)); + + row.skip(1).unwrap(); + row.write(header.get_candidate_name(id_worst_last)).unwrap(); + row.write(score_worst_last).unwrap(); + row.write(Percent(score_worst_last, total.into())).unwrap(); + for ((&score, &score_last), &running) in scores.iter().zip(scores_last.iter()).zip(running.iter()) { if running { let score_diff = score - score_last; - print!(",+{score_diff},+{:.3}%", score_diff / score_worst_last * 100.0); + row.write(score_diff).unwrap(); + row.write(PercentDiff(score_diff, score_worst_last)).unwrap(); } else { - print!(",,"); + row.skip(2).unwrap(); } } - println!(); + drop(row); + row = csv_out.row(); } - print!("{},,,", index + 1); + + row.write(index + 1).unwrap(); + row.skip(3).unwrap(); + for (&score, &running) in scores.iter().zip(running.iter()) { if running { - print!(",{score},{:.3}%", to_percent(score)); + row.write(score).unwrap(); + row.write(Percent(score, total.into())).unwrap(); } else { - print!(",,"); + row.skip(2).unwrap(); } } - println!(); + running[id_worst] = false; scores_last = Some(scores); id_worst_last = Some(id_worst); let mut total_votes = 0.0; - println!(); eprintln!("Round {}:", index+1); for (place, &(id, score)) in scores_ordered.iter().take(winner_count).enumerate() { - eprintln!(" {}. {}: {score}, {:.3}%", place + 1, names[id], to_percent(score)); + eprintln!(" {}. {}: {score}, {}", place + 1, header.get_candidate_name(id), Percent(score, total.into())); total_votes += score; } if scores_ordered.len() > winner_count { - eprintln!(" {}. {}: {score_worst}, {:.3}%", scores_ordered.len(), names[id_worst], to_percent(score_worst)); + eprintln!(" {}. {}: {score_worst}, {}", scores_ordered.len(), header.get_candidate_name(id_worst), Percent(score_worst, total.into())); } - eprintln!(" Total: {total_votes}, {:.3}%", to_percent(total_votes.into())); + eprintln!(" Total: {total_votes}, {}", Percent(total_votes.into(), total.into())); } } diff --git a/src/util.rs b/src/util.rs new file mode 100644 index 0000000..1b0a535 --- /dev/null +++ b/src/util.rs @@ -0,0 +1,7 @@ + +pub mod escape; +pub mod percent; + +pub use escape::{EscapeWriter, EscapeWriterOpts}; +pub use percent::Percent; + diff --git a/src/util/escape.rs b/src/util/escape.rs new file mode 100644 index 0000000..6a428f6 --- /dev/null +++ b/src/util/escape.rs @@ -0,0 +1,49 @@ + +pub struct EscapeWriterOpts<'a> { + pub chars: &'a [u8], + pub escape: u8, +} + +pub struct EscapeWriter<'a,T> { + opts: EscapeWriterOpts<'a>, + dst: T, +} + +impl<'a,T> EscapeWriter<'a,T> { + pub fn new(dst: T, opts: EscapeWriterOpts<'a>) -> Self { + Self { dst, opts } + } +} + +impl<'a> Default for EscapeWriterOpts<'a> { + fn default() -> Self { + Self { + escape: b'\\', + chars: b"\"", + } + } +} + +impl<'a,T> std::io::Write for EscapeWriter<'a,T> where T: std::io::Write { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let mut count = 0; + for &b in buf.iter() { + if b == self.opts.escape || self.opts.chars.contains(&b) { + if self.dst.write(&[self.opts.escape])? == 0 { + return Ok(count); + } + count += 1; + } + if self.dst.write(&[b])? == 0 { + return Ok(count); + } + count += 1; + } + Ok(count) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.dst.flush() + } +} + diff --git a/src/util/percent.rs b/src/util/percent.rs new file mode 100644 index 0000000..8a41ad7 --- /dev/null +++ b/src/util/percent.rs @@ -0,0 +1,21 @@ +use std::fmt::Display; + + +pub struct Percent(pub f64, pub f64); +pub struct PercentDiff(pub f64, pub f64); + +impl Display for Percent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:.3}%", self.0 / self.1 * 100.0) + } +} +impl Display for PercentDiff { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let v = self.0 / self.1 * 100.0; + write!(f, "{}{:.3}%", match v >= 0.0 { + true => '+', + false => '-', + }, v.abs()) + } +} +