/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021–2022 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, StageKind};
use crate::numbers::Number;

use itertools::Itertools;
use nohash_hasher::BuildNoHashHasher;

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
#[derive(Clone)]
pub struct BallotTree<'t, N: Number> {
	num_ballots: N,
	ballots: Vec<BallotInTree<'t, N>>,
	next_preferences: Option<Box<HashMap<&'t Candidate, BallotTree<'t, N>, BuildNoHashHasher<Candidate>>>>,
	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 [Candidate]) {
		let mut next_preferences: HashMap<&Candidate, BallotTree<N>, BuildNoHashHasher<Candidate>> = HashMap::with_capacity_and_hasher(candidates.len(), BuildNoHashHasher::default());
		let mut next_exhausted = BallotTree::new();
		
		for bit in self.ballots.iter() {
			if bit.up_to_pref < bit.ballot.preferences.len() {
				let preference = &bit.ballot.preferences[bit.up_to_pref];
				
				if preference.len() != 1 {
					todo!();
				}
				
				let candidate = &candidates[*preference.first().unwrap()];
				
				match next_preferences.get_mut(candidate) {
					Some(np_bt) => {
						np_bt.num_ballots += &bit.ballot.orig_value;
						np_bt.ballots.push(BallotInTree {
							ballot: bit.ballot,
							up_to_pref: bit.up_to_pref + 1,
						});
					}
					None => {
						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>, opts: &STVOptions)
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,
			up_to_pref: 0,
		});
		ballot_tree.num_ballots += &ballot.orig_value;
	}
	state.ballot_tree = Some(ballot_tree);
	
	// Distribute preferences
	distribute_preferences(state, opts);
	
	// Recalculate transfers
	for (_, count_card) in state.candidates.iter_mut() {
		count_card.transfers.assign(&count_card.votes);
	}
	state.exhausted.transfers.assign(&state.exhausted.votes);
	
	// Calculate loss by fraction - if minivoters used
	if let Some(orig_total) = &state.election.total_votes {
		let mut total_votes = state.candidates.values().fold(N::new(), |mut acc, cc| { acc += &cc.votes; acc });
		total_votes += &state.exhausted.votes;
		let lbf = orig_total - &total_votes;
		
		state.loss_fraction.votes = lbf.clone();
		state.loss_fraction.transfers = lbf;
	}
	
	state.title = StageKind::FirstPreferences;
	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>, opts: &STVOptions)
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, opts);
}

/// Distribute preferences recursively
///
/// Called by [distribute_preferences]
fn distribute_recursively<'t, N: Number>(candidates: &mut HashMap<&'t Candidate, CountCard<N>, BuildNoHashHasher<Candidate>>, exhausted: &mut CountCard<N>, tree: &mut BallotTree<'t, N>, remaining_multiplier: N, election: &'t Election<N>, opts: &STVOptions)
where
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
{
	// Descend tree if required
	if tree.next_exhausted.is_none() {
		tree.descend_tree(&election.candidates);
	}
	
	// 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
				let mut to_transfer = &remaining_multiplier * &cand_tree.num_ballots;
				if let Some(dps) = opts.round_votes {
					// NZ Meek STV rounds *up*!
					to_transfer.ceil_mut(dps);
				}
				count_card.votes += to_transfer;
			}
			CandidateState::Elected => {
				// Transfer according to elected candidate's keep value
				let mut to_transfer = &remaining_multiplier * &cand_tree.num_ballots * count_card.keep_value.as_ref().unwrap();
				if let Some(dps) = opts.round_votes {
					to_transfer.ceil_mut(dps);
				}
				count_card.votes += to_transfer;
				
				let mut new_remaining_multiplier = &remaining_multiplier * &(N::one() - count_card.keep_value.as_ref().unwrap());
				if let Some(dps) = opts.round_surplus_fractions {
					new_remaining_multiplier.ceil_mut(dps);
				}
				
				// Recurse
				distribute_recursively(candidates, exhausted, cand_tree, new_remaining_multiplier, election, opts);
			}
			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, opts);
			}
		}
	}
	
	// Credit exhausted votes at this level
	exhausted.votes += &remaining_multiplier * &tree.next_exhausted.as_ref().unwrap().as_ref().num_ballots;
}

fn recompute_keep_values<'s, N: Number>(state: &mut CountState<'s, N>, opts: &STVOptions, has_surplus: &[&'s Candidate]) {
	for candidate in has_surplus {
		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);
		
		if let Some(dps) = opts.round_values {
			// NZ Meek STV rounds *up*!
			count_card.keep_value.as_mut().unwrap().ceil_mut(dps);
		}
	}
}

/// Determine if the specified surpluses should be distributed, according to [STVOptions::meek_surplus_tolerance]
fn should_distribute_surpluses<'a, N: Number>(state: &CountState<'a, N>, has_surplus: &[&'a Candidate], opts: &STVOptions) -> bool
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
{
	if opts.meek_surplus_tolerance.ends_with('%') {
		// Distribute if any candidate has a surplus exceeding the tolerance
		let quota_tolerance = N::parse(&opts.meek_surplus_tolerance[0..opts.meek_surplus_tolerance.len()-1]) / N::from(100) + N::one();
		return has_surplus.iter().any(|c| {
			let count_card = &state.candidates[c];
			return &count_card.votes / state.quota.as_ref().unwrap() > quota_tolerance;
		});
	} else {
		// Distribute if the total surplus exceeds the tolerance
		let quota_tolerance = N::parse(&opts.meek_surplus_tolerance);
		let total_surpluses = has_surplus.iter()
			.fold(N::new(), |mut acc, c| { acc += &state.candidates[c].votes; acc -= state.quota.as_ref().unwrap(); acc });
		return total_surpluses > quota_tolerance;
	}
}

/// 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>,
{
	let quota = state.quota.as_ref().unwrap();
	let mut has_surplus: Vec<&Candidate> = state.election.candidates.iter()
		.filter(|c| {
			let count_card = &state.candidates[c];
			return count_card.state == CandidateState::Elected && &count_card.votes > quota;
		})
		.collect();
	
	let mut should_distribute = should_distribute_surpluses(state, &has_surplus, opts);
	if should_distribute {
		// Determine if surplues can be deferred
		if opts.defer_surpluses {
			let total_surpluses = has_surplus.iter()
				.fold(N::new(), |mut acc, c| { acc += &state.candidates[c].votes; acc -= quota; acc });
			if super::can_defer_surpluses(state, opts, &total_surpluses) {
				state.logger.log_literal(format!("Distribution of surpluses totalling {:.dps$} votes will be deferred.", total_surpluses, dps=opts.pp_decimals));
				return Ok(false);
			}
		}
		
		let mut surpluses_deferred = None; // Option<total_surpluses>
		let mut candidates_elected = None; // Option<LogEntry>
		
		let orig_candidates = state.candidates.clone();
		let orig_exhausted = state.exhausted.clone();
		
		let mut num_iterations: u32 = 0;
		
		while should_distribute {
			num_iterations += 1;
			
			// Recompute keep values
			recompute_keep_values(state, opts, &has_surplus);
			
			// Redistribute votes
			distribute_preferences(state, opts);
			
			// Recompute quota if more ballots have become exhausted
			super::calculate_quota(state, opts);
			
			if opts.immediate_elect {
				// Try to elect candidates
				if super::elect_hopefuls(state, opts, true)? {
					candidates_elected = Some(state.logger.entries.pop().unwrap());
					break;
				}
			}
			
			let quota = state.quota.as_ref().unwrap();
			has_surplus = state.election.candidates.iter()
				.filter(|c| {
					let count_card = &state.candidates[c];
					return count_card.state == CandidateState::Elected && &count_card.votes > quota;
				})
				.collect();
			
			should_distribute = should_distribute_surpluses(state, &has_surplus, opts);
			
			// Determine if surplues can be deferred
			if should_distribute && opts.defer_surpluses {
				let total_surpluses = has_surplus.iter()
					.fold(N::new(), |mut acc, c| { acc += &state.candidates[c].votes; acc -= quota; acc });
				if super::can_defer_surpluses(state, opts, &total_surpluses) {
					surpluses_deferred = Some(total_surpluses);
					break;
				}
			}
		}
		
		// Recalculate transfers
		let mut checksum = N::new();
		for (candidate, count_card) in state.candidates.iter_mut() {
			count_card.transfers = &count_card.votes - &orig_candidates[candidate].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.title = StageKind::SurplusesDistributed;
		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));
		}
		
		if let Some(total_surpluses) = surpluses_deferred {
			state.logger.log_literal(format!("Distribution of surpluses totalling {:.dps$} votes will be deferred.", total_surpluses, dps=opts.pp_decimals));
		}
		
		// If candidates were elected, retain that log entry
		if let Some(log_entry) = candidates_elected {
			state.logger.log(log_entry);
		}
		
		let kv_str = state.election.candidates.iter()
			.map(|c| (c, &state.candidates[c]))
			.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>,
{
	// NZ Meek STV: Iterate keep values one round before exclusion
	if opts.meek_nz_exclusion {
		let quota = state.quota.as_ref().unwrap();
		let has_surplus: Vec<&Candidate> = state.election.candidates.iter()
			.filter(|c| {
				let count_card = &state.candidates[c];
				return count_card.state == CandidateState::Elected && &count_card.votes > quota;
			})
			.collect();
		recompute_keep_values(state, opts, &has_surplus);
		
		let kv_str = state.election.candidates.iter()
			.map(|c| (c, &state.candidates[c]))
			.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));
	}
	
	// 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);
			count_card.finalised = true;
		}
	}
	
	let orig_candidates = state.candidates.clone();
	let orig_exhausted = state.exhausted.clone();
	
	distribute_preferences(state, opts);
	
	// Recalculate transfers
	let mut checksum = N::new();
	for (candidate, count_card) in state.candidates.iter_mut() {
		count_card.transfers = &count_card.votes - &orig_candidates[candidate].votes;
		checksum += &count_card.transfers;
	}
	state.exhausted.transfers = &state.exhausted.votes - &orig_exhausted.votes;
	checksum += &state.exhausted.transfers;
	state.loss_fraction.transfer(&-checksum);
}