/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021–2023 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::candmap::CandidateMap;
use crate::constraints::{Constraint, Constraints, ConstrainedGroup, ConstraintMatrix};
use crate::logger::Logger;
use crate::numbers::Number;
use crate::sharandom::SHARandom;
use crate::stv::{self, STVOptions};
use crate::stv::gregory::TransferTable;
use crate::stv::meek::BallotTree;

use itertools::Itertools;

#[cfg(not(target_arch = "wasm32"))]
use rkyv::{Archive, Deserialize, Serialize};
#[cfg(not(target_arch = "wasm32"))]
use crate::numbers::{SerializedNumber, SerializedOptionNumber};

use std::cmp::max;
use std::fmt;
use std::hash::{Hash, Hasher};

/// An election to be counted
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
#[derive(Clone)]
pub struct Election<N> {
	/// Name of the election
	pub name: String,
	/// Number of candidates to be elected
	pub seats: usize,
	/// [Vec] of [Candidate]s in the election
	pub candidates: Vec<Candidate>,
	/// Indexes of withdrawn candidates
	pub withdrawn_candidates: Vec<usize>,
	/// [Vec] of [Ballot]s cast in the election
	pub ballots: Vec<Ballot<N>>,
	/// Total value of [Ballot]s cast in the election
	///
	/// Used for [Election::realise_equal_rankings].
	#[cfg_attr(not(target_arch = "wasm32"), with(SerializedOptionNumber))]
	pub total_votes: Option<N>,
	/// Constraints on candidates
	pub constraints: Option<Constraints>,
}

impl<N: Number> Election<N> {
	/// Convert ballots with weight >1 to multiple ballots of weight 1
	///
	/// Assumes ballots have integer weight.
	pub fn normalise_ballots(&mut self) {
		let mut normalised_ballots = Vec::new();
		for ballot in self.ballots.iter() {
			let mut n = N::new();
			let one = N::one();
			while n < ballot.orig_value {
				let new_ballot = Ballot {
					orig_value: N::one(),
					preferences: ballot.preferences.clone(),
					has_equal_rankings: ballot.has_equal_rankings,
				};
				normalised_ballots.push(new_ballot);
				n += &one;
			}
		}
		self.ballots = normalised_ballots;
	}
	
	/// Convert ballots with equal rankings to strict-preference "minivoters"
	pub fn realise_equal_rankings(&mut self) {
		if !self.ballots.iter().any(|b| b.has_equal_rankings) {
			// No equal rankings
			return;
		}
		
		// Record total_votes so loss by fraction can be calculated
		// See crate::stv::gregory::distribute_first_preferences, etc.
		self.total_votes = Some(self.ballots.iter().fold(N::new(), |mut acc, b| { acc += &b.orig_value; acc }));
		
		let mut realised_ballots = Vec::with_capacity(self.ballots.len());
		for ballot in self.ballots.drain(..) {
			ballot.realise_equal_rankings_into(&mut realised_ballots);
		}
		self.ballots = realised_ballots;
	}
}

/// A candidate in an [Election]
#[derive(Clone, Eq)]
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
pub struct Candidate {
	/// Index of the candidate
	pub index: usize,
	/// Name of the candidate
	pub name: String,
	/// If this candidate is a dummy candidate (e.g. for --constraint-mode repeat_count)
	pub is_dummy: bool,
}

impl PartialEq for Candidate {
	// Custom implementation of eq for HashMap purposes, to improve performance
	//
	// SAFETY: Results in undefined behaviour if multiple Candidates are allowed to have the same index
	fn eq(&self, other: &Candidate) -> bool {
		return self.index == other.index;
	}
}

impl Hash for Candidate {
	fn hash<H: Hasher>(&self, hasher: &mut H) {
		// Custom implementation of hash for use with NoHashHasher, to improve performance
		hasher.write_usize(self.index);
	}
}
impl nohash_hasher::IsEnabled for Candidate {}

/// The current state of counting an [Election]
#[derive(Clone)]
pub struct CountState<'a, N: Number> {
	/// Pointer to the [Election] being counted
	pub election: &'a Election<N>,
	
	/// [CandidateMap] of [CountCard]s for each [Candidate] in the election
	pub candidates: CandidateMap<'a, CountCard<'a, N>>,
	/// [CountCard] representing the exhausted pile
	pub exhausted: CountCard<'a, N>,
	/// [CountCard] representing loss by fraction
	pub loss_fraction: CountCard<'a, N>,
	
	/// [BallotTree] for Meek STV
	pub ballot_tree: Option<BallotTree<'a, N>>,
	
	/// Values used to break ties, based on forwards tie-breaking
	pub forwards_tiebreak: Option<CandidateMap<'a, usize>>,
	/// Values used to break ties, based on backwards tie-breaking
	pub backwards_tiebreak: Option<CandidateMap<'a, usize>>,
	/// [SHARandom] for random tie-breaking
	pub random: Option<SHARandom<'a>>,
	
	/// Quota for election
	pub quota: Option<N>,
	/// Vote required for election
	///
	/// Only used in ERS97/ERS76.
	pub vote_required_election: Option<N>,
	
	/// Number of candidates who have been declared elected
	pub num_elected: usize,
	/// Number of candidates who have been declared excluded
	pub num_excluded: usize,
	
	/// [ConstraintMatrix] for constrained elections
	pub constraint_matrix: Option<ConstraintMatrix>,
	/// [RollbackState] when using [ConstraintMode::RepeatCount](crate::stv::ConstraintMode::RepeatCount)
	pub rollback_state: RollbackState<'a, N>,
	
	/// Transfer table for this surplus/exclusion
	pub transfer_table: Option<TransferTable<'a, N>>,
	
	/// The type of stage being counted, etc.
	pub title: StageKind<'a>,
	/// [Logger] for this stage of the count
	pub logger: Logger<'a>,
}

impl<'a, N: Number> CountState<'a, N> {
	/// Construct a new blank [CountState] for the given [Election]
	pub fn new(election: &'a Election<N>) -> Self {
		let mut state = CountState {
			election,
			candidates: CandidateMap::with_capacity(election.candidates.len()),
			exhausted: CountCard::new(),
			loss_fraction: CountCard::new(),
			ballot_tree: None,
			forwards_tiebreak: None,
			backwards_tiebreak: None,
			random: None,
			quota: None,
			vote_required_election: None,
			num_elected: 0,
			num_excluded: 0,
			constraint_matrix: None,
			rollback_state: RollbackState::Normal,
			transfer_table: None,
			title: StageKind::FirstPreferences,
			logger: Logger { entries: Vec::new() },
		};
		
		// Init candidate count cards
		for candidate in election.candidates.iter() {
			let mut count_card = CountCard::new();
			if candidate.is_dummy {
				count_card.state = CandidateState::Withdrawn;
			}
			state.candidates.insert(candidate, count_card);
		}
		
		// Set withdrawn candidates state
		for withdrawn_idx in election.withdrawn_candidates.iter() {
			state.candidates.get_mut(&election.candidates[*withdrawn_idx]).unwrap().state = CandidateState::Withdrawn;
		}
		
		// Init constraints
		if let Some(constraints) = &election.constraints {
			// Init constraint matrix
			let mut num_groups: Vec<usize> = constraints.0.iter().map(|c| c.groups.len()).collect();
			let mut cm = ConstraintMatrix::new(&mut num_groups[..]);
			
			// Init constraint matrix total cells min/max
			for (i, constraint) in constraints.0.iter().enumerate() {
				for (j, group) in constraint.groups.iter().enumerate() {
					let mut idx = vec![0; constraints.0.len()];
					idx[i] = j + 1;
					let mut cell = &mut cm[&idx];
					cell.min = group.min;
					cell.max = group.max;
				}
			}
			
			// Fill in grand total, etc.
			cm.update_from_state(state.election, &state.candidates);
			cm.init();
			//println!("{}", cm);
			
			// Require correct number of candidates to be elected
			let idx = vec![0; constraints.0.len()];
			cm[&idx].min = election.seats;
			cm[&idx].max = election.seats;
			
			state.constraint_matrix = Some(cm);
		}
		
		return state;
	}
	
	/// [Step](CountCard::step) every [CountCard] to prepare for the next stage
	pub fn step_all(&mut self) {
		for (_, count_card) in self.candidates.iter_mut() {
			count_card.step();
		}
		self.exhausted.step();
		self.loss_fraction.step();
	}
	
	/// List the candidates, and their current state, votes and transfers
	pub fn describe_candidates(&self, opts: &STVOptions) -> String {
		let mut candidates: Vec<(&Candidate, &CountCard<N>)>;
		
		if opts.sort_votes {
			// Sort by votes if requested
			candidates = self.candidates.iter().collect();
			// First sort by order of election (as a tie-breaker, if votes are equal)
			candidates.sort_unstable_by(|a, b| b.1.order_elected.cmp(&a.1.order_elected));
			// Then sort by votes
			candidates.sort_by(|a, b| a.1.votes.cmp(&b.1.votes));
			candidates.reverse();
		} else {
			candidates = self.election.candidates.iter()
				.map(|c| (c, &self.candidates[c]))
				.collect();
		}
		
		let mut result = String::new();
		
		for (candidate, count_card) in candidates {
			if candidate.is_dummy {
				continue;
			}
			
			match count_card.state {
				CandidateState::Hopeful => {
					result.push_str(&format!("- {}: {:.dps$} ({:.dps$})\n", candidate.name, count_card.votes, count_card.transfers, dps=opts.pp_decimals));
				}
				CandidateState::Guarded => {
					result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - Guarded\n", candidate.name, count_card.votes, count_card.transfers, dps=opts.pp_decimals));
				}
				CandidateState::Elected => {
					if let Some(kv) = &count_card.keep_value {
						result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - ELECTED {} (kv = {:.dps2$})\n", candidate.name, count_card.votes, count_card.transfers, count_card.order_elected, kv, dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
					} else {
						result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - ELECTED {}\n", candidate.name, count_card.votes, count_card.transfers, count_card.order_elected, dps=opts.pp_decimals));
					}
				}
				CandidateState::Doomed => {
					result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - Doomed\n", candidate.name, count_card.votes, count_card.transfers, dps=opts.pp_decimals));
				}
				CandidateState::Withdrawn => {
					if !opts.hide_excluded || !count_card.votes.is_zero() || !count_card.transfers.is_zero() {
						result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - Withdrawn\n", candidate.name, count_card.votes, count_card.transfers, dps=opts.pp_decimals));
					}
				}
				CandidateState::Excluded => {
					// If --hide-excluded, hide unless nonzero votes or nonzero transfers
					if !opts.hide_excluded || !count_card.votes.is_zero() || !count_card.transfers.is_zero() {
						result.push_str(&format!("- {}: {:.dps$} ({:.dps$}) - Excluded {}\n", candidate.name, count_card.votes, count_card.transfers, -count_card.order_elected, dps=opts.pp_decimals));
					}
				}
			}
		}
		return result;
	}
	
	/// Produce summary rows for the current stage
	pub fn describe_summary(&self, opts: &STVOptions) -> String {
		let mut result = String::new();
		
		result.push_str(&format!("Exhausted: {:.dps$} ({:.dps$})\n", self.exhausted.votes, self.exhausted.transfers, dps=opts.pp_decimals));
		result.push_str(&format!("Loss by fraction: {:.dps$} ({:.dps$})\n", self.loss_fraction.votes, self.loss_fraction.transfers, dps=opts.pp_decimals));
		
		let mut total_vote = self.candidates.iter().filter_map(|(c, cc)| if c.is_dummy { None } else { Some(cc) }).fold(N::new(), |mut acc, cc| { acc += &cc.votes; acc });
		total_vote += &self.exhausted.votes;
		total_vote += &self.loss_fraction.votes;
		result.push_str(&format!("Total votes: {:.dps$}\n", total_vote, dps=opts.pp_decimals));
		
		if self.election.seats == 1 {
			result.push_str(&format!("Majority: {:.dps$}\n", self.quota.as_ref().unwrap(), dps=opts.pp_decimals));
		} else {
			result.push_str(&format!("Quota: {:.dps$}\n", self.quota.as_ref().unwrap(), dps=opts.pp_decimals));
		}
		if stv::should_show_vre(opts) {
			if let Some(vre) = &self.vote_required_election {
				result.push_str(&format!("Vote required for election: {:.dps$}\n", vre, dps=opts.pp_decimals));
			}
		}
		
		return result;
	}
	
	// -------------------
	// HELPER CALCULATIONS
	
	/// Get the total vote, viz. the sum votes of all candidates
	pub fn total_vote(&self) -> N {
		return self.total_votes_of(self.candidates.values());
	}
	
	/// Get the total votes of the given candidates
	pub fn total_votes_of<I: Iterator<Item=&'a CountCard<'a, N>>>(&self, count_cards: I) -> N {
		return count_cards.fold(N::new(), |mut acc, cc| { acc += &cc.votes; acc });
	}
	
	/// Get the total active vote, viz. the votes of all continuing candidates plus all votes awaiting transfer
	pub fn active_vote(&self) -> N {
		return self.candidates.values().fold(N::new(), |mut acc, cc| {
			match cc.state {
				CandidateState::Elected => {
					if !cc.finalised && &cc.votes > self.quota.as_ref().unwrap() {
						acc += &cc.votes;
						acc -= self.quota.as_ref().unwrap();
					}
				}
				_ => {
					acc += &cc.votes;
				}
			}
			acc
		});
	}
	
	/// Get the total surplus
	/// 
	/// A candidate has a surplus if the candidate's votes are not finalised, and the votes exceed the quota.
	/// The votes of all other candidates are ignored.
	pub fn total_surplus(&self) -> N {
		return self.total_surplus_of(self.candidates.values());
	}
	
	/// Get the total surpluses of the given candidates
	/// 
	/// See [CountState::total_surplus].
	pub fn total_surplus_of<I: Iterator<Item=&'a CountCard<'a, N>>>(&self, count_cards: I) -> N {
		return count_cards.fold(N::new(), |mut acc, cc| {
			if !cc.finalised && &cc.votes > self.quota.as_ref().unwrap() {
				acc += &cc.votes;
				acc -= self.quota.as_ref().unwrap();
			}
			acc
		});
	}
}

/// The kind, title, etc. of the stage being counted
#[derive(Clone)]
pub enum StageKind<'a> {
	/// First preferences
	FirstPreferences,
	/// Surplus of ...
	SurplusOf(&'a Candidate),
	/// Exclusion of ...
	ExclusionOf(Vec<&'a Candidate>),
	/// Rolled back (--constraint-mode repeat_count)
	Rollback,
	/// Exhausted ballots (--constraint-mode repeat_count)
	RollbackExhausted,
	/// Ballots of ... (--constraint-mode repeat_count)
	BallotsOf(&'a Candidate),
	/// Surpluses distributed (Meek)
	SurplusesDistributed,
	/// Bulk election
	BulkElection,
}

impl<'a> StageKind<'a> {
	/// Return the "kind" portion of the title
	pub fn kind_as_string(&self) -> &'static str {
		return match self {
			StageKind::FirstPreferences => "",
			StageKind::SurplusOf(_) => "Surplus of",
			StageKind::ExclusionOf(_) => "Exclusion of",
			StageKind::Rollback => "",
			StageKind::RollbackExhausted => "",
			StageKind::BallotsOf(_) => "Ballots of",
			StageKind::SurplusesDistributed => "",
			StageKind::BulkElection => "",
		};
	}
}

impl<'a> fmt::Display for StageKind<'a> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		match self {
			StageKind::FirstPreferences => {
				return f.write_str("First preferences");
			}
			StageKind::SurplusOf(candidate) => {
				return f.write_fmt(format_args!("{} {}", self.kind_as_string(), candidate.name));
			}
			StageKind::ExclusionOf(candidates) => {
				return f.write_fmt(format_args!("{} {}", self.kind_as_string(), candidates.iter().map(|c| &c.name).sorted().join(", ")));
			}
			StageKind::Rollback => {
				return f.write_str("Constraints applied");
			}
			StageKind::RollbackExhausted => {
				return f.write_str("Exhausted ballots");
			}
			StageKind::BallotsOf(candidate) => {
				return f.write_fmt(format_args!("{} {}", self.kind_as_string(), candidate.name));
			}
			StageKind::SurplusesDistributed => {
				return f.write_str("Surpluses distributed");
			}
			StageKind::BulkElection => {
				return f.write_str("Bulk election");
			}
		}
	}
}

/// Current state of a [Candidate] during an election count
#[derive(Clone)]
pub struct CountCard<'a, N> {
	/// State of the candidate
	pub state: CandidateState,
	/// Order of election or exclusion
	///
	/// Positive integers represent order of election; negative integers represent order of exclusion.
	pub order_elected: isize,
	/// Whether distribution of this candidate's surpluses/transfer of excluded candidate's votes is complete
	pub finalised: bool,
	
	/// Net votes transferred to this candidate in this stage
	pub transfers: N,
	/// Votes of the candidate at the end of this stage
	pub votes: N,
	
	/// Net ballots transferred to this candidate in this stage
	pub ballot_transfers: N,
	
	/// Parcels of ballots assigned to this candidate
	pub parcels: Vec<Parcel<'a, N>>,
	
	/// Candidate's keep value (Meek STV)
	pub keep_value: Option<N>,
}

impl<'a, N: Number> CountCard<'a, N> {
	/// Returns a new blank [CountCard]
	pub fn new() -> Self {
		return CountCard {
			state: CandidateState::Hopeful,
			order_elected: 0,
			finalised: false,
			transfers: N::new(),
			votes: N::new(),
			ballot_transfers: N::new(),
			parcels: Vec::new(),
			keep_value: None,
		};
	}
	
	/// Transfer the given number of votes to this [CountCard], incrementing [transfers](CountCard::transfers) and [votes](CountCard::votes)
	pub fn transfer(&mut self, transfer: &'_ N) {
		self.transfers += transfer;
		self.votes += transfer;
	}
	
	/// Set [transfers](CountCard::transfers) to 0
	pub fn step(&mut self) {
		self.transfers = N::new();
		self.ballot_transfers = N::new();
	}
	
	/// Concatenate all parcels into a single parcel, leaving [parcels](CountCard::parcels) empty
	pub fn concat_parcels(&mut self) -> Vec<Vote<'a, N>> {
		let mut result = Vec::new();
		for parcel in self.parcels.iter_mut() {
			result.append(&mut parcel.votes);
		}
		return result;
	}
	
	/// Return the number of ballots across all parcels
	pub fn num_ballots(&self) -> N {
		return self.parcels.iter().fold(N::new(), |mut acc, p| { acc += p.num_ballots(); acc });
	}
}

/// Parcel of [Vote]s during a count
#[derive(Clone)]
pub struct Parcel<'a, N> {
	/// [Vote]s in this parcel
	pub votes: Vec<Vote<'a, N>>,
	/// Accumulated relative value of each [Vote] in this parcel
	pub value_fraction: N,
	/// Order for sorting with [crate::stv::ExclusionMethod::BySource]
	pub source_order: usize,
}

impl<'a, N: Number> Parcel<'a, N> {
	/// Return the number of ballots in this parcel
	pub fn num_ballots(&self) -> N {
		return self.votes.iter().fold(N::new(), |mut acc, v| { acc += &v.ballot.orig_value; acc });
	}
	
	/// Return the value of the votes in this parcel
	pub fn num_votes(&self) -> N {
		return self.num_ballots() * &self.value_fraction;
	}
}

/// Represents a [Ballot] with an associated value
#[derive(Clone)]
pub struct Vote<'a, N> {
	/// Ballot from which the vote is derived
	pub ballot: &'a Ballot<N>,
	/// Index of the next preference to examine
	pub up_to_pref: usize,
}

impl<'a, N> Vote<'a, N> {
	/// Get the next preference and increment `up_to_pref`
	///
	/// Assumes that each preference level contains only one preference.
	pub fn next_preference(&mut self) -> Option<usize> {
		if self.up_to_pref >= self.ballot.preferences.len() {
			return None;
		}
		
		let preference = &self.ballot.preferences[self.up_to_pref];
		self.up_to_pref += 1;
		
		return Some(*preference.first().unwrap());
	}
}

/// A record of a voter's preferences
#[derive(Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
pub struct Ballot<N> {
	/// Original value/weight of the ballot
	#[cfg_attr(not(target_arch = "wasm32"), with(SerializedNumber))]
	pub orig_value: N,
	/// Indexes of candidates preferenced at each level on the ballot
	pub preferences: Vec<Vec<usize>>,
	/// Whether multiple candidates are preferenced at the same level
	pub has_equal_rankings: bool,
}

impl<N: Number> Ballot<N> {
	/// Convert ballot with equal rankings to strict-preference "minivoters"
	pub fn realise_equal_rankings_into(self, dest: &mut Vec<Ballot<N>>) {
		if !self.has_equal_rankings {
			dest.push(self);
			return;
		}
		
		// Preferences for each minivoter
		let mut minivoters = vec![Vec::new()];
		
		for preference in self.preferences.iter() {
			if preference.len() == 1 {
				// Single preference so just add to the end of existing preferences
				for minivoter in minivoters.iter_mut() {
					minivoter.push(preference.clone());
				}
			} else {
				// Equal ranking
				// Get all possible permutations
				let permutations: Vec<Vec<usize>> = preference.iter().copied().permutations(preference.len()).collect();
				
				// Split into new "minivoters" for each possible permutation
				let mut new_minivoters = Vec::with_capacity(minivoters.len() * permutations.len());
				for permutation in permutations {
					for minivoter in minivoters.iter() {
						let mut new_minivoter = minivoter.clone();
						for p in permutation.iter() {
							new_minivoter.push(vec![*p]);
						}
						new_minivoters.push(new_minivoter);
					}
				}
				minivoters = new_minivoters;
			}
		}
		
		let weight_each = self.orig_value.clone() / N::from(minivoters.len());
		
		for minivoter in minivoters {
			dest.push(Ballot {
				orig_value: weight_each.clone(),
				preferences: minivoter,
				has_equal_rankings: false,
			});
		}
	}
}

/// State of a [Candidate] during a count
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum CandidateState {
	/// Hopeful (continuing candidate)
	Hopeful,
	/// Required by constraints to be guarded from exclusion
	Guarded,
	/// Declared elected
	Elected,
	/// Required by constraints to be doomed to be excluded
	Doomed,
	/// Withdrawn candidate
	Withdrawn,
	/// Declared excluded
	Excluded,
}

/// If --constraint-mode repeat_count and redistribution is required, tracks the ballot papers being redistributed
#[allow(missing_docs)]
#[derive(Clone)]
pub enum RollbackState<'a, N> {
	/// Not rolling back
	Normal,
	/// Start rolling back next stage
	NeedsRollback {
		candidates: Option<CandidateMap<'a, CountCard<'a, N>>>,
		exhausted: Option<CountCard<'a, N>>,
		constraint: &'a Constraint,
		group: &'a ConstrainedGroup
	},
	/// Rolling back
	RollingBack {
		candidates: Option<CandidateMap<'a, CountCard<'a, N>>>,
		exhausted: Option<CountCard<'a, N>>,
		candidate_distributing: Option<&'a Candidate>,
		constraint: Option<&'a Constraint>,
		group: Option<&'a ConstrainedGroup>
	},
}