/*  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 super::{STVError, STVOptions};

use crate::election::{Ballot, Candidate, CandidateState, CountCard, CountState, Election};
use crate::numbers::Number;

use itertools::Itertools;

use std::cmp::max;
use std::collections::HashMap;
use std::ops;

/// Ballot in a [BallotTree]
#[derive(Clone)]
struct BallotInTree<'b, N: Number> {
	ballot: &'b Ballot<N>,
	/// Index of the next preference to examine
	up_to_pref: usize,
}

/// Tree-packed ballot representation
pub struct BallotTree<'t, N: Number> {
	num_ballots: N,
	ballots: Vec<BallotInTree<'t, N>>,
	next_preferences: Option<Box<HashMap<&'t Candidate, BallotTree<'t, N>>>>,
	next_exhausted: Option<Box<BallotTree<'t, N>>>,
}

impl<'t, N: Number> BallotTree<'t, N> {
	/// Construct a new empty [BallotTree]
	fn new() -> Self {
		BallotTree {
			num_ballots: N::new(),
			ballots: Vec::new(),
			next_preferences: None,
			next_exhausted: None,
		}
	}
	
	/// Descend one level of the [BallotTree]
	fn descend_tree(&mut self, candidates: &'t Vec<Candidate>) {
		let mut next_preferences: HashMap<&Candidate, BallotTree<N>> = HashMap::new();
		let mut next_exhausted = BallotTree::new();
		
		for bit in self.ballots.iter() {
			if bit.up_to_pref < bit.ballot.preferences.len() {
				let candidate = &candidates[bit.ballot.preferences[bit.up_to_pref]];
				
				if next_preferences.contains_key(candidate) {
					let np_bt = next_preferences.get_mut(candidate).unwrap();
					np_bt.num_ballots += &bit.ballot.orig_value;
					np_bt.ballots.push(BallotInTree {
						ballot: bit.ballot,
						up_to_pref: bit.up_to_pref + 1,
					});
				} else {
					let mut np_bt = BallotTree::new();
					np_bt.num_ballots += &bit.ballot.orig_value;
					np_bt.ballots.push(BallotInTree {
						ballot: bit.ballot,
						up_to_pref: bit.up_to_pref + 1,
					});
					next_preferences.insert(candidate, np_bt);
				}
			} else {
				// Exhausted
				next_exhausted.num_ballots += &bit.ballot.orig_value;
				next_exhausted.ballots.push(bit.clone());
			}
		}
		
		self.next_preferences = Some(Box::new(next_preferences));
		self.next_exhausted = Some(Box::new(next_exhausted));
	}
}

/// Initialise keep values, ballot tree and distribute preferences
pub fn distribute_first_preferences<N: Number>(state: &mut CountState<N>)
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
{
	// Initialise keep values
	for (_, count_card) in state.candidates.iter_mut() {
		count_card.keep_value = Some(N::one());
	}
	
	// Initialise ballot tree
	let mut ballot_tree = BallotTree::new();
	for ballot in state.election.ballots.iter() {
		ballot_tree.ballots.push(BallotInTree {
			ballot: ballot,
			up_to_pref: 0,
		});
		ballot_tree.num_ballots += &ballot.orig_value;
	}
	state.ballot_tree = Some(ballot_tree);
	
	// Distribute preferences
	distribute_preferences(state);
	
	// Recalculate transfers
	for (_, count_card) in state.candidates.iter_mut() {
		count_card.transfers.assign(&count_card.votes);
	}
	state.exhausted.transfers.assign(&state.exhausted.votes);
	
	state.kind = None;
	state.title = "First preferences".to_string();
	state.logger.log_literal("First preferences distributed.".to_string());
}

/// (Re)distribute preferences according to candidate keep values
pub fn distribute_preferences<N: Number>(state: &mut CountState<N>)
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
{
	// Reset votes
	for (_, count_card) in state.candidates.iter_mut() {
		//count_card.orig_votes = N::new();
		//count_card.transfers = N::new();
		count_card.votes = N::new();
	}
	state.exhausted.votes = N::new();
	
	distribute_recursively(&mut state.candidates, &mut state.exhausted, state.ballot_tree.as_mut().unwrap(), N::one(), &state.election);
}

/// Distribute preferences recursively
///
/// Called by [distribute_preferences]
fn distribute_recursively<'t, N: Number>(candidates: &mut HashMap<&'t Candidate, CountCard<N>>, exhausted: &mut CountCard<N>, tree: &mut BallotTree<'t, N>, remaining_multiplier: N, election: &'t Election<N>)
where
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
{
	// Descend tree if required
	if let None = tree.next_exhausted {
		tree.descend_tree(&election.candidates);
	}
	
	// FIXME: Possibility of infinite loop if malformed inputs?
	
	// TODO: Round transfers?
	
	// Credit votes at this level
	for (candidate, cand_tree) in tree.next_preferences.as_mut().unwrap().as_mut().iter_mut() {
		let count_card = candidates.get_mut(candidate).unwrap();
		match count_card.state {
			CandidateState::Hopeful | CandidateState::Guarded | CandidateState::Doomed => {
				// Hopeful candidate has keep value 1, so transfer entire remaining value
				count_card.votes += &remaining_multiplier * &cand_tree.num_ballots;
			}
			CandidateState::Elected => {
				// Transfer according to elected candidate's keep value
				count_card.votes += &remaining_multiplier * &cand_tree.num_ballots * count_card.keep_value.as_ref().unwrap();
				let new_remaining_multiplier = &remaining_multiplier * &(N::one() - count_card.keep_value.as_ref().unwrap());
				// Recurse
				distribute_recursively(candidates, exhausted, cand_tree, new_remaining_multiplier, election);
			}
			CandidateState::Excluded | CandidateState::Withdrawn => {
				// Excluded candidate has keep value 0, so skip over this candidate
				// Recurse
				distribute_recursively(candidates, exhausted, cand_tree, remaining_multiplier.clone(), election);
			}
		}
	}
	
	// Credit exhausted votes at this level
	exhausted.votes += &remaining_multiplier * &tree.next_exhausted.as_ref().unwrap().as_ref().num_ballots;
}

/// Recalculate all candidate keep factors to distribute all surpluses according to the Meek method
pub fn distribute_surpluses<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> Result<bool, STVError>
where
	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>,
{
	// TODO: Make configurable
	let quota_tolerance = N::one() / N::from(100000) + N::one();
	
	let quota = state.quota.as_ref().unwrap();
	let mut has_surplus: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of tie
		.filter(|c| {
			let count_card = state.candidates.get(c).unwrap();
			return count_card.state == CandidateState::Elected && (&count_card.votes / quota > quota_tolerance);
		})
		.collect();
	
	if !has_surplus.is_empty() {
		// TODO: Defer surpluses?
		
		let orig_candidates = state.candidates.clone();
		let orig_exhausted = state.exhausted.clone();
		
		let mut num_iterations: u32 = 0;
		
		while !has_surplus.is_empty() {
			num_iterations += 1;
			
			// Recompute keep values
			for candidate in has_surplus.into_iter() {
				let count_card = state.candidates.get_mut(candidate).unwrap();
				count_card.keep_value = Some(count_card.keep_value.take().unwrap() * state.quota.as_ref().unwrap() / &count_card.votes);
			}
			
			// Redistribute votes
			distribute_preferences(state);
			
			// Recompute quota if more ballots have become exhausted
			super::calculate_quota(state, opts);
			
			//println!("Debug {}", num_iterations);
			
			let quota = state.quota.as_ref().unwrap();
			has_surplus = state.election.candidates.iter()
				.filter(|c| {
					let count_card = state.candidates.get(c).unwrap();
					return count_card.state == CandidateState::Elected && (&count_card.votes / quota > quota_tolerance);
				})
				.collect();
		}
		
		// Recalculate transfers
		let mut checksum = N::new();
		for (candidate, count_card) in state.candidates.iter_mut() {
			count_card.transfers = &count_card.votes - &orig_candidates.get(candidate).unwrap().votes;
			checksum += &count_card.transfers;
		}
		state.exhausted.transfers = &state.exhausted.votes - &orig_exhausted.votes;
		checksum += &state.exhausted.transfers;
		state.loss_fraction.transfer(&-checksum);
		
		// Remove intermediate logs on quota calculation
		state.logger.entries.clear();
		
		state.kind = None;
		state.title = "Surpluses distributed".to_string();
		if num_iterations == 1 {
			state.logger.log_literal("Surpluses distributed, requiring 1 iteration.".to_string());
		} else {
			state.logger.log_literal(format!("Surpluses distributed, requiring {} iterations.", num_iterations));
		}
		
		let kv_str = state.election.candidates.iter()
			.map(|c| (c, state.candidates.get(c).unwrap()))
			.filter(|(_, cc)| cc.state == CandidateState::Elected)
			.sorted_unstable_by(|a, b| a.1.order_elected.cmp(&b.1.order_elected))
			.map(|(c, cc)| format!("{} ({:.dps2$})", c.name, cc.keep_value.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2)))
			.join(", ");
		
		state.logger.log_literal(format!("Keep values of elected candidates are: {}.", kv_str));
		
		return Ok(true);
	}
	return Ok(false);
}

/// Exclude the given candidates according to the Meek method
pub fn exclude_candidates<'a, N: Number>(state: &mut CountState<'a, N>, _opts: &STVOptions, excluded_candidates: Vec<&'a Candidate>)
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
{
	// Used to give bulk excluded candidate the same order_elected
	let order_excluded = state.num_excluded + 1;
	
	for excluded_candidate in excluded_candidates.into_iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		
		if count_card.state != CandidateState::Excluded {
			count_card.state = CandidateState::Excluded;
			state.num_excluded += 1;
			count_card.order_elected = -(order_excluded as isize);
		}
	}
	
	let orig_candidates = state.candidates.clone();
	let orig_exhausted = state.exhausted.clone();
	
	distribute_preferences(state);
	
	// Recalculate transfers
	let mut checksum = N::new();
	for (candidate, count_card) in state.candidates.iter_mut() {
		count_card.transfers = &count_card.votes - &orig_candidates.get(candidate).unwrap().votes;
		checksum += &count_card.transfers;
	}
	state.exhausted.transfers = &state.exhausted.votes - &orig_exhausted.votes;
	checksum += &state.exhausted.transfers;
	state.loss_fraction.transfer(&-checksum);
}