/*  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::{ExclusionMethod, NextPreferencesEntry, NextPreferencesResult, STVError, STVOptions, SumSurplusTransfersMode, SurplusMethod, SurplusOrder};

use crate::constraints;
use crate::election::{Candidate, CandidateState, CountState, Parcel, Vote};
use crate::numbers::Number;
use crate::ties;

use itertools::Itertools;

use std::cmp::max;
use std::ops;

/// Distribute first preference votes according to the Gregory method
pub fn distribute_first_preferences<N: Number>(state: &mut CountState<N>) {
	let votes = state.election.ballots.iter().map(|b| Vote {
		ballot: b,
		value: b.orig_value.clone(),
		up_to_pref: 0,
	}).collect();
	
	let result = super::next_preferences(state, votes);
	
	// Transfer candidate votes
	for (candidate, entry) in result.candidates.into_iter() {
		let parcel = entry.votes as Parcel<N>;
		let count_card = state.candidates.get_mut(candidate).unwrap();
		count_card.parcels.push(parcel);
		count_card.transfer(&entry.num_votes);
	}
	
	// Transfer exhausted votes
	let parcel = result.exhausted.votes as Parcel<N>;
	state.exhausted.parcels.push(parcel);
	state.exhausted.transfer(&result.exhausted.num_votes);
	
	state.kind = None;
	state.title = "First preferences".to_string();
	state.logger.log_literal("First preferences distributed.".to_string());
}

/// Distribute the largest surplus according to the Gregory method, based on [STVOptions::surplus]
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::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	let quota = state.quota.as_ref().unwrap();
	let has_surplus: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of tie
		.filter(|c| {
			let cc = &state.candidates[c];
			&cc.votes > quota && cc.parcels.iter().any(|p| !p.is_empty())
		})
		.collect();
	
	if !has_surplus.is_empty() {
		let total_surpluses = has_surplus.iter()
			.fold(N::new(), |acc, c| acc + &state.candidates[c].votes - quota);
		
		// Determine if surplues can be deferred
		if opts.defer_surpluses {
			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);
			}
		}
		
		// Distribute top candidate's surplus
		let max_cands = match opts.surplus_order {
			SurplusOrder::BySize => {
				ties::multiple_max_by(&has_surplus, |c| &state.candidates[c].votes)
			}
			SurplusOrder::ByOrder => {
				ties::multiple_min_by(&has_surplus, |c| state.candidates[c].order_elected)
			}
		};
		let elected_candidate = if max_cands.len() > 1 {
			super::choose_highest(state, opts, max_cands)?
		} else {
			max_cands[0]
		};
		
		distribute_surplus(state, &opts, elected_candidate);
		
		return Ok(true);
	}
	return Ok(false);
}

/// Return the denominator of the transfer value
fn calculate_surplus_denom<N: Number>(surplus: &N, result: &NextPreferencesResult<N>, transferable_votes: &N, weighted: bool, transferable_only: bool) -> Option<N>
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>
{
	if transferable_only {
		let total_units = if weighted { &result.total_votes } else { &result.total_ballots };
		let exhausted_units = if weighted { &result.exhausted.num_votes } else { &result.exhausted.num_ballots };
		let transferable_units = total_units - exhausted_units;
		
		if transferable_votes > surplus {
			return Some(transferable_units);
		} else {
			return None;
		}
	} else {
		if weighted {
			return Some(result.total_votes.clone());
		} else {
			return Some(result.total_ballots.clone());
		}
	}
}

/// Return the reweighted value of the vote after being transferred
fn reweight_vote<N: Number>(
	num_votes: &N,
	num_ballots: &N,
	surplus: &N,
	weighted: bool,
	surplus_fraction: &Option<N>,
	surplus_denom: &Option<N>,
	round_tvs: Option<usize>,
	rounding: Option<usize>) -> N
{
	let mut result;
	
	match surplus_denom {
		Some(v) => {
			if let Some(_) = round_tvs {
				// Rounding requested: use the rounded transfer value
				if weighted {
					result = num_votes.clone() * surplus_fraction.as_ref().unwrap();
				} else {
					result = num_ballots.clone() * surplus_fraction.as_ref().unwrap();
				}
			} else {
				// Avoid unnecessary rounding error by first multiplying by the surplus
				if weighted {
					result = num_votes.clone() * surplus / v;
				} else {
					result = num_ballots.clone() * surplus / v;
				}
			}
		}
		None => {
			result = num_votes.clone();
		}
	}
	
	// Round down if requested
	if let Some(dps) = rounding {
		result.floor_mut(dps);
	}
	
	return result;
}

/// Compute the number of votes to credit to a continuing candidate during a surplus transfer, based on [STVOptions::sum_surplus_transfers]
fn sum_surplus_transfers<N: Number>(entry: &NextPreferencesEntry<N>, surplus: &N, is_weighted: bool, surplus_fraction: &Option<N>, surplus_denom: &Option<N>, _state: &mut CountState<N>, opts: &STVOptions) -> N
where
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
{
	match opts.sum_surplus_transfers {
		SumSurplusTransfersMode::SingleStep => {
			// Calculate transfer across all votes
			//state.logger.log_literal(format!("Transferring {:.0} ballot papers, totalling {:.dps$} votes.", entry.num_ballots, entry.num_votes, dps=opts.pp_decimals));
			return reweight_vote(&entry.num_votes, &entry.num_ballots, surplus, is_weighted, surplus_fraction, surplus_denom, opts.round_tvs, opts.round_votes);
		}
		SumSurplusTransfersMode::ByValue => {
			// Sum transfers by value
			let mut result = N::new();
			
			// Sort into parcels by value
			let mut votes: Vec<&Vote<N>> = entry.votes.iter().collect();
			votes.sort_unstable_by(|a, b| (&a.value / &a.ballot.orig_value).cmp(&(&b.value / &b.ballot.orig_value)));
			for (_value, parcel) in &votes.into_iter().group_by(|v| &v.value / &v.ballot.orig_value) {
				let mut num_votes = N::new();
				let mut num_ballots = N::new();
				for vote in parcel {
					num_votes += &vote.value;
					num_ballots += &vote.ballot.orig_value;
				}
				//state.logger.log_literal(format!("Transferring {:.0} ballot papers, totalling {:.dps$} votes, received at value {:.dps2$}.", num_ballots, num_votes, value, dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
				result += reweight_vote(&num_votes, &num_ballots, surplus, is_weighted, surplus_fraction, surplus_denom, opts.round_tvs, opts.round_votes);
			}
			
			return result;
		}
		SumSurplusTransfersMode::PerBallot => {
			// Sum transfer per each individual ballot
			// TODO: This could be moved to distribute_surplus to avoid looping over the votes and calculating transfer values twice
			let mut result = N::new();
			for vote in entry.votes.iter() {
				result += reweight_vote(&vote.value, &vote.ballot.orig_value, surplus, is_weighted, surplus_fraction, surplus_denom, opts.round_tvs, opts.round_votes);
			}
			//state.logger.log_literal(format!("Transferring {:.0} ballot papers, totalling {:.dps$} votes.", entry.num_ballots, entry.num_votes, dps=opts.pp_decimals));
			return result;
		}
	}
}

/// Distribute the surplus of a given candidate according to the Gregory method, based on [STVOptions::surplus]
fn distribute_surplus<N: Number>(state: &mut CountState<N>, opts: &STVOptions, elected_candidate: &Candidate)
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	state.logger.log_literal(format!("Surplus of {} distributed.", elected_candidate.name));
	
	let count_card = &state.candidates[elected_candidate];
	let surplus = &count_card.votes - state.quota.as_ref().unwrap();
	
	let votes;
	match opts.surplus {
		SurplusMethod::WIG | SurplusMethod::UIG => {
			// Inclusive Gregory
			votes = count_card.parcels.concat();
		}
		SurplusMethod::EG => {
			// Exclusive Gregory
			// Should be safe to unwrap() - or else how did we get a quota!
			votes = state.candidates.get_mut(elected_candidate).unwrap().parcels.pop().unwrap();
		}
		_ => { panic!("Invalid --surplus for Gregory method"); }
	}
	
	// Count next preferences
	let result = super::next_preferences(state, votes);
	
	state.kind = Some("Surplus of");
	state.title = String::from(&elected_candidate.name);
	
	// Transfer candidate votes
	// TODO: Refactor??
	let is_weighted = match opts.surplus {
		SurplusMethod::WIG => { true }
		SurplusMethod::UIG | SurplusMethod::EG => { false }
		SurplusMethod::Meek => { todo!() }
	};
	
	let transferable_votes = &result.total_votes - &result.exhausted.num_votes;
	let surplus_denom = calculate_surplus_denom(&surplus, &result, &transferable_votes, is_weighted, opts.transferable_only);
	let mut surplus_fraction;
	match surplus_denom {
		Some(ref v) => {
			surplus_fraction = Some(surplus.clone() / v);
			
			// Round down if requested
			if let Some(dps) = opts.round_tvs {
				surplus_fraction.as_mut().unwrap().floor_mut(dps);
			}
			
			if opts.transferable_only {
				if &result.total_ballots - &result.exhausted.num_ballots == N::one() {
					state.logger.log_literal(format!("Transferring 1 transferable ballot, totalling {:.dps$} transferable votes, with surplus fraction {:.dps2$}.", transferable_votes, surplus_fraction.as_ref().unwrap(), dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
				} else {
					state.logger.log_literal(format!("Transferring {:.0} transferable ballots, totalling {:.dps$} transferable votes, with surplus fraction {:.dps2$}.", &result.total_ballots - &result.exhausted.num_ballots, transferable_votes, surplus_fraction.as_ref().unwrap(), dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
				}
			} else {
				if result.total_ballots == N::one() {
					state.logger.log_literal(format!("Transferring 1 ballot, totalling {:.dps$} votes, with surplus fraction {:.dps2$}.", result.total_votes, surplus_fraction.as_ref().unwrap(), dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
				} else {
					state.logger.log_literal(format!("Transferring {:.0} ballots, totalling {:.dps$} votes, with surplus fraction {:.dps2$}.", result.total_ballots, result.total_votes, surplus_fraction.as_ref().unwrap(), dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
				}
			}
		}
		None => {
			surplus_fraction = None;
			
			if opts.transferable_only {
				state.logger.log_literal(format!("Transferring {:.0} transferable ballots, totalling {:.dps$} transferable votes, at values received.", &result.total_ballots - &result.exhausted.num_ballots, transferable_votes, dps=opts.pp_decimals));
			} else {
				state.logger.log_literal(format!("Transferring {:.0} ballots, totalling {:.dps$} votes, at values received.", result.total_ballots, result.total_votes, dps=opts.pp_decimals));
			}
		}
	}
	
	let mut checksum = N::new();
	
	for (candidate, entry) in result.candidates.into_iter() {
		// Credit transferred votes
		let candidate_transfers = sum_surplus_transfers(&entry, &surplus, is_weighted, &surplus_fraction, &surplus_denom, state, opts);
		let count_card = state.candidates.get_mut(candidate).unwrap();
		count_card.transfer(&candidate_transfers);
		checksum += candidate_transfers;
		
		let mut parcel = entry.votes as Parcel<N>;
		
		// Reweight votes
		for vote in parcel.iter_mut() {
			vote.value = reweight_vote(&vote.value, &vote.ballot.orig_value, &surplus, is_weighted, &surplus_fraction, &surplus_denom, opts.round_tvs, opts.round_weights);
		}
		
		count_card.parcels.push(parcel);
	}
	
	// Credit exhausted votes
	let mut exhausted_transfers;
	if opts.transferable_only {
		if transferable_votes > surplus {
			// No ballots exhaust
			exhausted_transfers = N::new();
		} else {
			exhausted_transfers = &surplus - &transferable_votes;
			
			if let Some(dps) = opts.round_votes {
				exhausted_transfers.floor_mut(dps);
			}
		}
	} else {
		exhausted_transfers = sum_surplus_transfers(&result.exhausted, &surplus, is_weighted, &surplus_fraction, &surplus_denom, state, opts);
	}
	
	state.exhausted.transfer(&exhausted_transfers);
	checksum += exhausted_transfers;
	
	// Transfer exhausted votes
	let parcel = result.exhausted.votes as Parcel<N>;
	state.exhausted.parcels.push(parcel);
	
	// Finalise candidate votes
	let count_card = state.candidates.get_mut(elected_candidate).unwrap();
	count_card.transfers = -&surplus;
	count_card.votes.assign(state.quota.as_ref().unwrap());
	checksum -= surplus;
	
	count_card.parcels.clear(); // Mark surpluses as done
	
	// Update loss by fraction
	state.loss_fraction.transfer(&-checksum);
}

/// Perform one stage of a candidate exclusion according to the Gregory method, based on [STVOptions::exclusion]
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::Div<&'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.iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		
		// Rust borrow checker is unhappy if we try to put this in exclude_hopefuls ??!
		if count_card.state != CandidateState::Excluded {
			count_card.state = CandidateState::Excluded;
			state.num_excluded += 1;
			count_card.order_elected = -(order_excluded as isize);
			
			constraints::update_constraints(state, opts);
		}
	}
	
	// Determine votes to transfer in this stage
	let mut votes = Vec::new();
	let mut votes_remain;
	let mut checksum = N::new();
	
	match opts.exclusion {
		ExclusionMethod::SingleStage => {
			// Exclude in one round
			for excluded_candidate in excluded_candidates.iter() {
				let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
				votes.append(&mut count_card.parcels.concat());
				count_card.parcels.clear();
				
				// Update votes
				let votes_transferred = votes.iter().fold(N::new(), |acc, v| acc + &v.value);
				checksum -= &votes_transferred;
				count_card.transfer(&-votes_transferred);
			}
			votes_remain = false;
		}
		ExclusionMethod::ByValue => {
			// Exclude by value
			let excluded_with_votes: Vec<&&Candidate> = excluded_candidates.iter().filter(|c| !state.candidates[*c].parcels.is_empty()).collect();
			
			if excluded_with_votes.is_empty() {
				votes_remain = false;
			} else {
				// If candidates to exclude still having votes, select only those with the greatest value
				let max_value = excluded_with_votes.iter()
					.map(|c| state.candidates[*c].parcels.iter()
						.map(|p| p.iter().map(|v| &v.value / &v.ballot.orig_value).max().unwrap())
						.max().unwrap())
					.max().unwrap();
				
				votes_remain = false;
				
				for excluded_candidate in excluded_with_votes.iter() {
					let count_card = state.candidates.get_mut(*excluded_candidate).unwrap();
					
					// Filter out just those votes with max_value
					let mut remaining_votes = Vec::new();
					
					let cand_votes = count_card.parcels.concat();
					
					let mut votes_transferred = N::new();
					for vote in cand_votes.into_iter() {
						if &vote.value / &vote.ballot.orig_value == max_value {
							votes_transferred += &vote.value;
							votes.push(vote);
						} else {
							remaining_votes.push(vote);
						}
					}
					
					if !remaining_votes.is_empty() {
						votes_remain = true;
					}
					
					// Leave remaining votes with candidate (as one parcel)
					count_card.parcels = vec![remaining_votes];
					
					// Update votes
					checksum -= &votes_transferred;
					count_card.transfer(&-votes_transferred);
				}
			}
		}
		ExclusionMethod::ParcelsByOrder => {
			// Exclude by parcel by order
			if excluded_candidates.len() > 1 {
				panic!("--exclusion parcels_by_order is incompatible with --bulk-exclude");
			}
			
			let count_card = state.candidates.get_mut(excluded_candidates[0]).unwrap();
			
			if count_card.parcels.is_empty() {
				votes_remain = false;
			} else {
				votes = count_card.parcels.remove(0);
				votes_remain = !count_card.parcels.is_empty();
				
				// Update votes
				let votes_transferred = votes.iter().fold(N::new(), |acc, v| acc + &v.value);
				checksum -= &votes_transferred;
				count_card.transfer(&-votes_transferred);
			}
		}
		_ => panic!()
	}
	
	if !votes.is_empty() {
		let value = &votes[0].value / &votes[0].ballot.orig_value;
		
		// Count next preferences
		let result = super::next_preferences(state, votes);
		
		if let ExclusionMethod::SingleStage = opts.exclusion {
			if result.total_ballots == N::one() {
				state.logger.log_literal(format!("Transferring 1 ballot, totalling {:.dps$} votes.", result.total_votes, dps=opts.pp_decimals));
			} else {
				state.logger.log_literal(format!("Transferring {:.0} ballots, totalling {:.dps$} votes.", result.total_ballots, result.total_votes, dps=opts.pp_decimals));
			}
		} else {
			if result.total_ballots == N::one() {
				state.logger.log_literal(format!("Transferring 1 ballot, totalling {:.dps$} votes, received at value {:.dps2$}.", result.total_votes, value, dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
			} else {
				state.logger.log_literal(format!("Transferring {:.0} ballots, totalling {:.dps$} votes, received at value {:.dps2$}.", result.total_ballots, result.total_votes, value, dps=opts.pp_decimals, dps2=max(opts.pp_decimals, 2)));
			}
		}
		
		// Transfer candidate votes
		for (candidate, entry) in result.candidates.into_iter() {
			let parcel = entry.votes as Parcel<N>;
			let count_card = state.candidates.get_mut(candidate).unwrap();
			count_card.parcels.push(parcel);
			
			// Round transfers
			let mut candidate_transfers = entry.num_votes;
			if let Some(dps) = opts.round_votes {
				candidate_transfers.floor_mut(dps);
			}
			count_card.transfer(&candidate_transfers);
			checksum += candidate_transfers;
		}
		
		// Transfer exhausted votes
		let parcel = result.exhausted.votes as Parcel<N>;
		state.exhausted.parcels.push(parcel);
		
		let mut exhausted_transfers = result.exhausted.num_votes;
		if let Some(dps) = opts.round_votes {
			exhausted_transfers.floor_mut(dps);
		}
		state.exhausted.transfer(&exhausted_transfers);
		checksum += exhausted_transfers;
	}
	
	if !votes_remain {
		// Finalise candidate votes
		for excluded_candidate in excluded_candidates.into_iter() {
			let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
			checksum -= &count_card.votes;
			count_card.transfers -= &count_card.votes;
			count_card.votes = N::new();
		}
		
		if let ExclusionMethod::SingleStage = opts.exclusion {
		} else {
			state.logger.log_literal("Exclusion complete.".to_string());
		}
	}
	
	// Update loss by fraction
	state.loss_fraction.transfer(&-checksum);
}

/// Perform one stage of a candidate exclusion according to the Wright method
pub fn wright_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>,
	for<'r> &'r N: ops::Div<&'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.iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		
		// Rust borrow checker is unhappy if we try to put this in exclude_hopefuls ??!
		if count_card.state != CandidateState::Excluded {
			count_card.state = CandidateState::Excluded;
			state.num_excluded += 1;
			count_card.order_elected = -(order_excluded as isize);
		}
		
		constraints::update_constraints(state, opts);
	}
	
	// Reset count
	for (_, count_card) in state.candidates.iter_mut() {
		if count_card.order_elected > 0 {
			count_card.order_elected = 0;
		}
		count_card.parcels.clear();
		count_card.votes = N::new();
		count_card.transfers = N::new();
		count_card.state = match count_card.state {
			CandidateState::Withdrawn => CandidateState::Withdrawn,
			CandidateState::Excluded => CandidateState::Excluded,
			_ => CandidateState::Hopeful,
		};
	}
	
	state.exhausted.votes = N::new();
	state.exhausted.transfers = N::new();
	state.loss_fraction.votes = N::new();
	state.loss_fraction.transfers = N::new();
	
	state.num_elected = 0;
	
	let orig_title = state.title.clone();
	
	// Redistribute first preferences
	super::distribute_first_preferences(state, opts);
	
	state.kind = Some("Exclusion of");
	state.title = orig_title;
	
	// Trigger recalculation of quota within stv::count_one_stage
	state.quota = None;
	state.vote_required_election = None;
}
