OpenTally/src/parser/csp.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 crate::election::{Ballot, Candidate, Election};
use crate::numbers::Number;

use csv::ReaderBuilder;
use itertools::Itertools;

use std::collections::HashMap;
use std::io::Read;

/// Parse the given CSP file
pub fn parse_reader<R: Read, N: Number>(reader: R) -> Election<N> {
	// Read CSV file
	let mut reader = ReaderBuilder::new()
		.has_headers(true)
		.from_reader(reader);
	
	// Read candidates
	let mut candidates = Vec::new();
	let mut idx_map = HashMap::new(); // Map csp index -> candidates index
	
	for (i, cand_name) in reader.headers().expect("Syntax Error").into_iter().enumerate() {
		if cand_name == "$mult" {
			continue;
		}
		
		idx_map.insert(i, candidates.len());
		candidates.push(Candidate {
			name: cand_name.to_string(),
		});
	}
	
	// Read ballots
	let mut ballots = Vec::new();
	
	for record in reader.into_records() {
		let record = record.expect("Syntax Error");
		
		let mut value = N::one();
		
		// Record preferences
		let mut preferences = Vec::new(); // Vec of (ranking, candidate index)
		for (csv_index, preference) in record.into_iter().enumerate() {
			match idx_map.get(&csv_index) {
				Some(cand_index) => {
					// Preference
					if preference.len() == 0 || preference == "-" {
						continue;
					}
					
					let preference: usize = preference.parse().expect("Syntax Error");
					if preference == 0 {
						continue;
					}
					
					preferences.push((preference, cand_index));
				}
				None => {
					// $mult column
					let mult: usize = preference.parse().expect("Syntax Error");
					if mult == 1 {
						continue;
					}
					value = N::from(mult);
				}
			}
		}
		
		// Sort by ranking
		let mut unique_rankings: Vec<usize> = preferences.iter().map(|(r, _)| *r).unique().collect();
		unique_rankings.sort();
		
		let mut sorted_preferences = Vec::with_capacity(preferences.len());
		for ranking in unique_rankings {
			// Filter for preferences at this ranking
			let prefs_this_ranking: Vec<usize> = preferences.iter()
				.filter_map(|(r, i)| if *r == ranking { Some(**i) } else { None })
				.collect();
			sorted_preferences.push(prefs_this_ranking);
		}
		
		ballots.push(Ballot {
			orig_value: value,
			preferences: sorted_preferences,
		});
	}
	
	return Election {
		name: String::new(),
		seats: 0,
		candidates: candidates,
		withdrawn_candidates: Vec::new(),
		ballots: ballots,
		total_votes: None,
		constraints: None,
	};
}