wybr 0.0.6

//! VoteTree tally for STV and Borda elections

use crate::tally::Metadata;
use crate::tally::Tally;
use std::collections::HashMap;
use std::collections::HashSet;

#[derive(Copy, Clone)]
///Vote transferring type
pub enum Transfer {
    ///Meek transfer; each rank gets given candidate's weight times the remaining vote part
    Meek,
    ///Warren transfer; each ranks gets given candidate's weight or the whole remaining vote,
    ///whatever is smaller
    Warren,
}

struct StvBranch {
    pub count: u64,
    children: HashMap<u32, StvBranch>,
}
impl StvBranch {
    fn new() -> StvBranch {
        StvBranch {
            count: 0,
            children: HashMap::new(),
        }
    }
    fn add(&mut self, vote: &[u32], count: u64) -> u64 {
        self.count += count;
        if vote.is_empty() {
            self.count
        } else {
            self.children
                .entry(vote[0])
                .or_insert(StvBranch {
                    count: 0,
                    children: HashMap::new(),
                })
                .add(&vote[1..], count)
        }
    }
    fn stream_ballots<F: FnMut(u64, &[u32])>(&self, reader: &mut F, stub: &mut Vec<u32>) -> u64 {
        //TODO: Elegant, then in order of sorted keys in c?
        let mut exhausted = 0;
        for (&c, deeper) in &self.children {
            stub.push(c);
            exhausted += deeper.stream_ballots(reader, stub);
            stub.pop();
        }

        if self.count > exhausted {
            if stub.is_empty() {
                panic!("Len 0 but not exhausted count??");
            }
            //This has short ballots
            reader(self.count - exhausted, stub);
        }
        self.count
    }
    fn distribute_votes(&self, scores: &mut HashMap<u32, u64>, eliminated: &HashSet<u32>) -> u64 {
        let mut assigned: u64 = 0;
        for (&cand, deeper) in &self.children {
            if !eliminated.contains(&cand) {
                //Cand gets all the vote
                *scores.entry(cand).or_insert(0) += deeper.count;
                assigned += deeper.count;
            } else {
                //Cand gets nothing, vote moves further
                assigned += deeper.distribute_votes(scores, eliminated);
            }
        }
        assigned
    }
    fn transfer_votes(
        &self,
        scores: &mut HashMap<u32, u64>,
        weights: &HashMap<u32, u64>,
        vote: u64,
        base: u64,
        transfer: Transfer,
    ) -> u64 {
        use std::cmp::min;
        //Vote is moved down according to weight
        let mut assigned: u64 = 0;
        for (&c, deeper) in &self.children {
            let given = match transfer {
                //c gets its weight * remaining part of vote
                Transfer::Meek => (vote * weights.get(&c).unwrap_or(&0)) / base,
                //c gets its weight or the remaining vote, whatever is smaller
                Transfer::Warren => min(vote, *weights.get(&c).unwrap_or(&0)),
            };
            if given > 0 {
                *scores.entry(c).or_insert(0) += deeper.count * given;
                assigned += deeper.count * given;
            }
            if given < vote {
                assigned += deeper.transfer_votes(scores, weights, vote - given, base, transfer);
            }
        }
        assigned
    }
    fn count_ranks(
        &self,
        points: &mut HashMap<(u32, u32), u64>,
        skipped: &HashSet<u32>,
        depth: u32,
    ) {
        for (&c, deeper) in &self.children {
            if !skipped.contains(&c) {
                *points.entry((c, depth)).or_insert(0) += deeper.count;
                deeper.count_ranks(points, skipped, depth + 1);
            } else {
                //Skip, hence go deeper without increasing depth
                deeper.count_ranks(points, skipped, depth);
            }
        }
    }
}

pub struct VoteTree {
    candidates: u32,
    root: StvBranch,
    meta: Metadata,
}
impl VoteTree {
    pub fn new(candidates: u32) -> VoteTree {
        VoteTree {
            candidates,
            root: StvBranch::new(),
            meta: Metadata::empty(),
        }
    }
    pub fn get_candidates(&self) -> u32 {
        self.candidates
    }
    pub fn ballot_count(&self) -> u64 {
        self.root.count
    }
    pub fn stream_ballots<F: FnMut(u64, &[u32])>(&self, reader: &mut F) {
        let mut vote: Vec<u32> = Vec::new();
        self.root.stream_ballots(reader, &mut vote);
    }
    pub fn assign_votes(&self, eliminated: &HashSet<u32>) -> (u64, HashMap<u32, u64>) {
        let mut scores = HashMap::new();
        let excess = self.root.count - self.root.distribute_votes(&mut scores, eliminated);
        (excess, scores)
    }
    pub fn transfer_votes_fp(
        &self,
        weights: &HashMap<u32, u64>,
        base: u64,
        transfer: Transfer,
    ) -> (u64, HashMap<u32, u64>) {
        //Ensure no overflow is possible with this base
        // in fact, meek/warren routines shall check this and produce some error.
        assert!(u64::MAX / base > self.root.count);
        let mut scores = HashMap::new();
        let total = self.root.count * base;
        let excess = total
            - self
                .root
                .transfer_votes(&mut scores, weights, base, base, transfer);
        (excess, scores)
    }
    #[cfg(test)]
    fn first_vote_count(&self, candidate: u32) -> u64 {
        if let Some(b) = self.root.children.get(&candidate) {
            b.count
        } else {
            0
        }
    }
    pub fn count_ranks(&self, skipped: &HashSet<u32>) -> HashMap<(u32, u32), u64> {
        let mut points: HashMap<(u32, u32), u64> = HashMap::new();
        self.root.count_ranks(&mut points, skipped, 0);
        points
    }
}

impl Tally for VoteTree {
    fn get_meta(&self) -> &'_ Metadata {
        &self.meta
    }
    fn get_candidates(&self) -> u32 {
        self.candidates
    }
}

use std::iter::FromIterator;
///Iterator collector for fast&dirty creation of VoteTrees in tests and examples
impl<'a> FromIterator<&'a (u64, Vec<u32>)> for VoteTree {
    fn from_iter<I: IntoIterator<Item = &'a (u64, Vec<u32>)>>(i: I) -> Self {
        let mut root = StvBranch::new();
        let mut maxc = 0;
        for (rep, votes) in i {
            root.add(votes, *rep);
            maxc = std::cmp::max(maxc, *votes.iter().max().unwrap_or(&0));
        }
        VoteTree {
            root,
            candidates: maxc + 1,
            meta: Metadata::empty(),
        }
    }
}

use self::VoteToken::*;
use crate::tally::{VoteReadError, VoteToken};
///Iterator collector for creating `VoteTree` from a stream of `VoteToken`s, as returned by ballot
///file parsers or user-defined sources like DB cursors, etc.
impl FromIterator<VoteToken> for Result<VoteTree, VoteReadError> {
    fn from_iter<I: IntoIterator<Item = VoteToken>>(i: I) -> Self {
        //Global state
        let mut withdrawn: HashSet<u32> = HashSet::new();
        let mut candidates: Option<u32> = None;
        let mut candidate_names: HashMap<u32, String> = HashMap::new();
        let mut seats: Option<u32> = None;
        let mut root: StvBranch = StvBranch::new();
        let mut title: Option<String> = None;

        //Per-ballot state
        use std::collections::BTreeMap;
        let mut sorted_ballot: BTreeMap<i32, u32> = BTreeMap::new();
        let mut ballot_repeat: Option<u64> = None;

        for token in i {
            match token {
                ReadFailure(x) => return Err(x),
                DeclareCandidates(x) => candidates = candidates.or(Some(x)),
                DeclareSeats(x) => seats = seats.or(Some(x)),
                WithdrawCandidate(x) => {
                    withdrawn.insert(x);
                }
                ElectionTitle(x) => title = title.or(Some(x)),
                BallotRepeat(x) => ballot_repeat = Some(x),
                Vote(pref, cand) => {
                    if sorted_ballot.insert(pref, cand).is_some() {
                        //Duplicated vote, escape!
                        return Err(VoteReadError::EqualVote);
                    }
                }
                EndBallot => {
                    let x = sorted_ballot.values().cloned().collect::<Vec<u32>>();
                    root.add(&x, ballot_repeat.unwrap());
                    ballot_repeat = None;
                    sorted_ballot.clear();
                }
                CandidateName(cand, name) => {
                    candidate_names.insert(cand, name);
                }
                _ => (),
            }
        }
        if let Some(candidates) = candidates {
            Ok(VoteTree {
                root,
                candidates,
                meta: Metadata::new(seats, withdrawn, candidate_names, title),
            })
        } else {
            Err(VoteReadError::InvalidMeta)
        }
    }
}

#[cfg(test)]
mod tests {
    use self::Transfer::*;
    use super::*;
    #[test]
    fn stv_branch_test() {
        let mut x = StvBranch::new();
        x.add(&[3, 5, 1, 7, 2], 4);
        let n = x.add(&[3, 5, 1, 7, 2], 3);
        assert_eq!(n, 7);
        assert_eq!(x.add(&[], 0), 7);
        assert_eq!(x.add(&[3], 0), 7);
        assert_eq!(x.add(&[4], 0), 0);
    }
    #[test]
    fn transfer_votes() {
        let x: VoteTree = [
            (1, vec![0, 1]),
            (2, vec![0, 2]),
            (3, vec![1, 0]),
            (4, vec![1, 2]),
            (5, vec![2, 0]),
            (6, vec![2, 1]),
        ]
        .iter()
        .collect();
        let base = 1_000_000;

        for &transfer in [Meek, Warren].iter() {
            let weights_one: HashMap<u32, u64> = (0..3).map(|i| (i as u32, base)).collect();
            let (excess, score) = x.transfer_votes_fp(&weights_one, base, transfer);
            assert_eq!(excess, 0);
            for i in 0..3 {
                assert_eq!(score[&i], x.first_vote_count(i) * base);
            }
        }

        let weights_half: HashMap<u32, u64> = [(0, base / 2), (1, base / 2), (2, base / 2)]
            .iter()
            .map(|(a, b)| (*a as u32, *b as u64))
            .collect();
        //Meek
        let (excess, score) = x.transfer_votes_fp(&weights_half, base, Meek);
        //Will be perfectly accurate provided that base is divisible by 4
        assert_eq!(score[&0], 3 * base / 2 + (3 + 5) * base / 4);
        assert_eq!(score[&1], 7 * base / 2 + (1 + 6) * base / 4);
        assert_eq!(score[&2], 11 * base / 2 + (4 + 2) * base / 4);
        assert_eq!(excess, (base * (1 + 2 + 3 + 4 + 5 + 6)) / 4);
        assert_eq!(
            score.iter().map(|(_, v)| *v).sum::<u64>() + excess,
            x.ballot_count() * base
        );

        //Warren
        let (excess, score) = x.transfer_votes_fp(&weights_half, base, Warren);
        //Will be perfectly accurate provided that base is divisible by 2
        assert_eq!(score[&0], 3 * base / 2 + (3 + 5) * base / 2);
        assert_eq!(score[&1], 7 * base / 2 + (1 + 6) * base / 2);
        assert_eq!(score[&2], 11 * base / 2 + (4 + 2) * base / 2);
        assert_eq!(excess, 0);
        assert_eq!(
            score.iter().map(|(_, v)| *v).sum::<u64>() + excess,
            x.ballot_count() * base
        );
    }
    #[test]
    fn vote_tree_collect() {
        let inpt = [(1, vec![0, 1]), (2, vec![0, 2]), (3, vec![1, 7, 3])];
        let x: VoteTree = inpt.iter().collect();
        assert_eq!(x.get_candidates(), 8);
        assert_eq!(x.ballot_count(), 6);
    }
    #[test]
    fn transfer_zero_weight() {
        let x: VoteTree = [
            (3, vec![0, 2, 3]),
            (4, vec![0, 2, 1]),
            (2, vec![3, 0, 2]),
            (1, vec![1]),
            (2, vec![1, 3, 2, 0]),
            (1, vec![2, 3, 1]),
        ]
        .iter()
        .collect();
        let base = 1_000_000;

        let vsum = x.root.count;

        let mut w: HashMap<u32, u64> = (0..4).map(|i| (i as u32, base)).collect();
        w.remove(&1);
        for &transfer in [Meek, Warren].iter() {
            let (excess, score) = x.transfer_votes_fp(&w, base, transfer);
            assert_eq!(excess + score.values().sum::<u64>(), vsum * base);
        }
    }
    #[test]
    fn ballot_streaming() {
        let x: VoteTree = [
            (1, vec![0]),
            (2, vec![0, 1]),
            (3, vec![0, 1, 2]),
            (4, vec![2]),
            (5, vec![2, 1]),
            (6, vec![2, 1, 0]),
        ]
        .iter()
        .collect();
        x.stream_ballots(&mut |rep, vec| match rep {
            1 => assert_eq!(vec, &[0]),
            2 => assert_eq!(vec, &[0, 1]),
            3 => assert_eq!(vec, &[0, 1, 2]),
            4 => assert_eq!(vec, &[2]),
            5 => assert_eq!(vec, &[2, 1]),
            6 => assert_eq!(vec, &[2, 1, 0]),
            _ => panic!("Invalid ballot generated"),
        });
    }
    #[test]
    fn count_ranks() {
        let x: VoteTree = [
            (1, vec![0]),
            (2, vec![0, 1]),
            (3, vec![0, 1, 2]),
            (4, vec![2]),
            (5, vec![2, 1]),
            (6, vec![2, 1, 0]),
        ]
        .iter()
        .collect();
        let ans = x.count_ranks(&HashSet::new());
        assert_eq!(ans[&(0, 0)], 6);
        assert_eq!(ans[&(0, 2)], 6);
        assert_eq!(ans[&(1, 1)], 16);
        assert_eq!(ans[&(2, 0)], 15);
        assert_eq!(ans[&(2, 2)], 3);
    }
    #[test]
    fn count_ranks_skip() {
        let x: VoteTree = [
            (1, vec![0]),
            (2, vec![0, 1]),
            (3, vec![0, 1, 2]),
            (4, vec![2]),
            (5, vec![2, 1]),
            (6, vec![2, 1, 0]),
        ]
        .iter()
        .collect();
        let ans = x.count_ranks(&([0].iter().cloned().collect()));
        assert_eq!(ans[&(1, 1)], 11);
        assert_eq!(ans[&(2, 1)], 3);
        assert_eq!(ans[&(2, 0)], 15);
        assert_eq!(ans[&(1, 0)], 5);
        let ans2 = x.count_ranks(&([0, 2].iter().cloned().collect()));
        assert_eq!(ans2[&(1, 0)], 16);
    }
    #[test]
    fn assign_votes_discrete() {
        let x: VoteTree = [
            (1, vec![0, 1]),
            (2, vec![0, 2]),
            (3, vec![1, 0]),
            (4, vec![1, 2]),
            (5, vec![2, 0]),
            (6, vec![2, 1]),
        ]
        .iter()
        .collect();

        let empty_hs = HashSet::new();
        let (excess, score) = x.assign_votes(&empty_hs);
        assert_eq!(excess, 0);
        assert_eq!(
            score,
            (0..3).map(|e| (e as u32, x.first_vote_count(e))).collect()
        );

        let (excess2, score2) = x.assign_votes(&([0u32, 2u32].iter().cloned().collect()));
        assert_eq!(*score2.get(&0).unwrap_or(&0), 0);
        assert_eq!(*score2.get(&2).unwrap_or(&0), 0);
        assert_eq!(score2[&1], x.ballot_count() - excess2);
    }
    #[test]
    fn assign_zero() {
        let x: VoteTree = [
            (3, vec![0, 2, 3]),
            (4, vec![0, 2, 1]),
            (2, vec![3, 0, 2]),
            (1, vec![1]),
            (2, vec![1, 3, 2, 0]),
            (1, vec![2, 3, 1]),
        ]
        .iter()
        .collect();

        let vsum = x.root.count;
        let (excess, score) = x.assign_votes(&([1u32].iter().cloned().collect()));
        assert_eq!(excess + score.values().sum::<u64>(), vsum);
    }
    #[test]
    fn stv_from_blt() {
        use crate::io::BltReader;
        use std::io::Cursor;
        const EASY_BLT: &'static [u8] =
            b" 3  2\n -3 \n 1 1 2 0\n2 1 3 0\n3 2 1 0 \n 4 0 0 2 3 0 0\n 5 3 1 0\n6 3 2 0\n 0 \n \"A\"\n  \n\"B\" \n\"C\"\n \"Title\" \n";
        let mut cursor = Cursor::new(EASY_BLT);
        let blt_reader = BltReader::new(&mut cursor);
        let vc = blt_reader
            .collect::<Result<VoteTree, VoteReadError>>()
            .unwrap();

        let (excess, score) = vc.assign_votes(&HashSet::new());
        assert_eq!(score[&0], 3);
        assert_eq!(score[&1], 7);
        assert_eq!(score[&2], 11);
        assert_eq!(excess, 0);
        assert_eq!(vc.get_candidates(), 3);
        assert_eq!(vc.get_seats(), 2);
    }
}