/*  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 crate::constraints;
use crate::election::{Candidate, CandidateState, CountState, Parcel, StageKind, Vote};
use crate::numbers::Number;
use crate::stv::{STVOptions, SampleMethod, SurplusMethod};

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

use super::{NextPreferencesResult, STVError};

/// Return the denominator of the surplus fraction
///
/// Returns `None` if transferable ballots <= surplus (i.e. all transferable ballots are transferred at full value).
fn calculate_surplus_denom<N: Number>(surplus: &N, result: &NextPreferencesResult<N>, transferable_ballots: &N, transferable_only: bool) -> Option<N>
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>
{
	if transferable_only {
		if transferable_ballots > surplus {
			return Some(transferable_ballots.clone());
		} else {
			return None;
		}
	} else {
		return Some(result.total_ballots.clone());
	}
}

/// Distribute the surplus of a given candidate according to the random subset method, based on [STVOptions::surplus]
pub fn distribute_surplus<'a, N: Number>(state: &mut CountState<'a, N>, opts: &STVOptions, elected_candidate: &'a Candidate) -> Result<(), 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>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	state.title = StageKind::SurplusOf(&elected_candidate);
	state.logger.log_literal(format!("Surplus of {} distributed.", elected_candidate.name));
	
	let count_card = state.candidates.get_mut(elected_candidate).unwrap();
	let surplus = &count_card.votes - state.quota.as_ref().unwrap();
	
	let mut votes;
	match opts.surplus {
		SurplusMethod::Cincinnati => {
			// Inclusive
			votes = count_card.concat_parcels();
		}
		SurplusMethod::Hare => {
			// Exclusive
			// Should be safe to unwrap() - or else how did we get a quota!
			votes = count_card.parcels.pop().unwrap().votes;
		}
		_ => unreachable!()
	}
	
	match opts.sample {
		SampleMethod::StratifyLR | SampleMethod::StratifyFloor => {
			// Stratified by next available preference (round fractions according to largest remainders)
			let result = super::next_preferences(state, votes);
			
			let transferable_ballots = &result.total_ballots - &result.exhausted.num_ballots;
			let surplus_denom = calculate_surplus_denom(&surplus, &result, &transferable_ballots, opts.transferable_only);
			let mut surplus_fraction;
			match &surplus_denom {
				Some(v) => {
					surplus_fraction = Some(surplus.clone() / v);
					
					// Round down if requested
					if let Some(dps) = opts.round_surplus_fractions {
						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!("Examining 1 transferable ballot, with surplus fraction {:.dps2$}.", surplus_fraction.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2)));
						} else {
							state.logger.log_literal(format!("Examining {:.0} transferable ballots, with surplus fraction {:.dps2$}.", &result.total_ballots - &result.exhausted.num_ballots, surplus_fraction.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2)));
						}
					} else {
						if result.total_ballots == N::one() {
							state.logger.log_literal(format!("Examining 1 ballot, with surplus fraction {:.dps2$}.", surplus_fraction.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2)));
						} else {
							state.logger.log_literal(format!("Examining {:.0} ballots, with surplus fraction {:.dps2$}.", result.total_ballots, surplus_fraction.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2)));
						}
					}
				}
				None => {
					surplus_fraction = None;
					
					// This can only happen if --transferable-only
					if result.total_ballots == N::one() {
						state.logger.log_literal("Transferring 1 ballot at full value.".to_string());
					} else {
						state.logger.log_literal(format!("Transferring {:.0} ballots at full value.", result.total_ballots));
					}
				}
			}
			
			let mut candidate_transfers_remainders: HashMap<Option<&Candidate>, (N, N)> = HashMap::new(); // None -> exhausted pile
			
			for (candidate, entry) in result.candidates.iter() {
				// Calculate votes to transfer
				let mut candidate_transfers;
				let remainder;
				match surplus_fraction {
					Some(ref f) => {
						match opts.sample {
							SampleMethod::StratifyLR => {
								// Incompatible with --round-surplus-fractions
								candidate_transfers = &entry.num_ballots * &surplus / surplus_denom.as_ref().unwrap();
								candidate_transfers.floor_mut(0);
								remainder = (&entry.num_ballots * &surplus) % surplus_denom.as_ref().unwrap();
							}
							SampleMethod::StratifyFloor => {
								match opts.round_surplus_fractions {
									Some(_) => {
										candidate_transfers = &entry.num_ballots * f;
									}
									None => {
										candidate_transfers = &entry.num_ballots * &surplus / surplus_denom.as_ref().unwrap();
									}
								}
								candidate_transfers.floor_mut(0);
								remainder = N::new();
							}
							_ => unreachable!()
						}
					}
					None => {
						// All ballots transferred
						candidate_transfers = entry.num_ballots.clone();
						remainder = N::new();
					}
				}
				
				candidate_transfers_remainders.insert(Some(candidate), (candidate_transfers, remainder));
			}
			
			// Calculate exhausted votes to transfer
			if !opts.transferable_only {
				let mut exhausted_transfers;
				let remainder;
				match opts.sample {
					SampleMethod::StratifyLR => {
						// Incompatible with --round-surplus-fractions
						exhausted_transfers = &result.exhausted.num_ballots * &surplus / surplus_denom.as_ref().unwrap();
						exhausted_transfers.floor_mut(0);
						remainder = (&result.exhausted.num_ballots * &surplus) % surplus_denom.as_ref().unwrap();
					}
					SampleMethod::StratifyFloor => {
						match opts.round_surplus_fractions {
							Some(_) => {
								exhausted_transfers = &result.exhausted.num_ballots * surplus_fraction.as_ref().unwrap();
							}
							None => {
								exhausted_transfers = &result.exhausted.num_ballots * &surplus / surplus_denom.as_ref().unwrap();
							}
						}
						exhausted_transfers.floor_mut(0);
						remainder = N::new();
					}
					_ => unreachable!()
				}
				
				candidate_transfers_remainders.insert(None, (exhausted_transfers, remainder));
			}
			
			if opts.sample == SampleMethod::StratifyLR {
				// Round remainders to remove loss by fraction
				let transferred = candidate_transfers_remainders.values().fold(N::new(), |acc, (t, _)| acc + t);
				let loss_fraction = &surplus - &transferred;
				if !loss_fraction.is_zero() && surplus_fraction.is_some() {
					let n_to_round: usize = format!("{:.0}", loss_fraction).parse().expect("Loss by fraction overflows usize");
					
					let mut cands_by_remainder = candidate_transfers_remainders.keys().cloned().collect::<Vec<_>>();
					
					// Sort by whole parts
					// Compare b to a to sort high-low
					cands_by_remainder.sort_unstable_by(|a, b| candidate_transfers_remainders[b].0.cmp(&candidate_transfers_remainders[a].0));
					
					// Then sort by remainders
					cands_by_remainder.sort_by(|a, b| candidate_transfers_remainders[b].1.cmp(&candidate_transfers_remainders[a].1));
					
					// Select top remainders
					let mut top_remainders = cands_by_remainder.iter().take(n_to_round).collect::<Vec<_>>();
					
					// Check for tied remainders
					if candidate_transfers_remainders[top_remainders.last().unwrap()] == candidate_transfers_remainders[cands_by_remainder.iter().nth(n_to_round).unwrap()] {
						// Get the top entry
						let top_entry = &candidate_transfers_remainders[top_remainders.last().unwrap()];
						
						// Separate out tied entries
						top_remainders = top_remainders.into_iter().filter(|c| &candidate_transfers_remainders[c] != top_entry).collect();
						let mut tied_top = cands_by_remainder.iter()
							.filter_map(|c| if let Some(c2) = c { if &candidate_transfers_remainders[c] == top_entry { Some(*c2) } else { None } } else { None })
							.collect::<Vec<_>>();
						
						// Get top entries by tie-breaking method
						for _ in 0..(n_to_round-top_remainders.len()) {
							let cand = super::choose_highest(state, opts, &tied_top, "Which fraction to round up?")?;
							tied_top.remove(tied_top.iter().position(|c| *c == cand).unwrap());
							top_remainders.push(cands_by_remainder.iter().find(|c| **c == Some(cand)).unwrap());
						}
					}
					
					// Round up top remainders
					for candidate in top_remainders {
						candidate_transfers_remainders.get_mut(candidate).unwrap().0 += N::one();
					}
				}
			}
			
			let mut checksum = N::new();
			
			for (candidate, entry) in result.candidates.into_iter() {
				// Credit transferred votes
				let candidate_transfers = &candidate_transfers_remainders[&Some(candidate)].0;
				let candidate_transfers_usize: usize = format!("{:.0}", candidate_transfers).parse().expect("Transfer overflows usize");
				
				let count_card = state.candidates.get_mut(candidate).unwrap();
				count_card.transfer(&candidate_transfers);
				checksum += candidate_transfers;
				
				let parcel = Parcel {
					votes: entry.votes.into_iter().rev().take(candidate_transfers_usize).rev().collect(),
					value_fraction: N::one(),
					source_order: state.num_elected + state.num_excluded,
				};
				
				count_card.parcels.push(parcel);
			}
			
			// Credit exhausted votes
			let mut exhausted_transfers;
			if opts.transferable_only {
				if transferable_ballots > surplus {
					// No ballots exhaust
					exhausted_transfers = N::new();
				} else {
					exhausted_transfers = &surplus - &transferable_ballots;
					if let Some(dps) = opts.round_votes {
						exhausted_transfers.floor_mut(dps);
					}
				}
			} else {
				exhausted_transfers = candidate_transfers_remainders.remove(&None).unwrap().0;
			}
			let exhausted_transfers_usize: usize = format!("{:.0}", exhausted_transfers).parse().expect("Transfer overflows usize");
			
			state.exhausted.transfer(&exhausted_transfers);
			checksum += exhausted_transfers;
			
			// Transfer exhausted votes
			let parcel = Parcel {
				votes: result.exhausted.votes.into_iter().rev().take(exhausted_transfers_usize).rev().collect(),
				value_fraction: N::one(),
				source_order: state.num_elected + state.num_excluded,
			};
			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;
			
			// Update loss by fraction
			state.loss_fraction.transfer(&-checksum);
		}
		SampleMethod::ByOrder => {
			// Ballots by order
			// FIXME: This is untested
			
			state.logger.log_literal(format!("Examining {:.0} ballots.", votes.len())); // votes.len() is total ballots as --normalise-ballots is required
			
			// Transfer candidate votes
			while &state.candidates[elected_candidate].votes > state.quota.as_ref().unwrap() {
				match votes.pop() {
					Some(vote) => {
						// Transfer to next preference
						transfer_ballot(state, opts, elected_candidate, vote, opts.transferable_only)?;
						if state.num_elected == state.election.seats {
							return Ok(());
						}
					}
					None => {
						// We have run out of ballot papers
						// Remaining ballot papers exhaust
						
						let surplus = &state.candidates[elected_candidate].votes - state.quota.as_ref().unwrap();
						state.exhausted.transfer(&surplus);
						state.candidates.get_mut(elected_candidate).unwrap().transfer(&-surplus);
						break;
					}
				}
			}
		}
		SampleMethod::NthBallot => {
			// Every nth-ballot (Cincinnati-style)
			
			// Calculate skip value
			let total_ballots = votes.len();
			let mut skip_fraction = N::from(total_ballots) / &surplus;
			skip_fraction.round_mut(0);
			
			state.logger.log_literal(format!("Examining {:.0} ballots, with skip value {:.0}.", total_ballots, skip_fraction));
			
			// Number the votes
			let mut numbered_votes: HashMap<usize, Vote<N>> = HashMap::new();
			for (i, vote) in votes.into_iter().enumerate() {
				numbered_votes.insert(i, vote);
			}
			
			// Transfer candidate votes
			let skip_value: usize = format!("{:.0}", skip_fraction).parse().expect("Skip value overflows usize");
			let mut iteration = 0;
			let mut index = skip_value - 1; // Subtract 1 as votes are 0-indexed
			
			while &state.candidates[elected_candidate].votes > state.quota.as_ref().unwrap() {
				// Transfer one vote to next available preference
				match numbered_votes.remove(&index) {
					Some(vote) => {
						transfer_ballot(state, opts, elected_candidate, vote, opts.transferable_only)?;
						if state.num_elected == state.election.seats {
							return Ok(());
						}
					}
					None => {
						// We have run out of ballot papers
						// Remaining ballot papers exhaust
						let surplus = &state.candidates[elected_candidate].votes - state.quota.as_ref().unwrap();
						state.exhausted.transfer(&surplus);
						state.candidates.get_mut(elected_candidate).unwrap().transfer(&-surplus);
						break;
					}
				}
				
				index += skip_value;
				if index >= total_ballots {
					iteration += 1;
					index = iteration + skip_value - 1;
				}
			}
		}
	}
	
	let count_card = state.candidates.get_mut(elected_candidate).unwrap();
	count_card.finalised = true; // Mark surpluses as done
	
	return Ok(());
}

/// Transfer the given ballot paper to its next available preference, and check for candidates meeting the quota if --sample-per-ballot
///
/// If `ignore_nontransferable`, does nothing if --transferable-only and the ballot is nontransferable.
fn transfer_ballot<'a, N: Number>(state: &mut CountState<'a, N>, opts: &STVOptions, source_candidate: &Candidate, mut vote: Vote<'a, N>, ignore_nontransferable: bool) -> Result<(), STVError> {
	// Get next preference
	let mut next_candidate = None;
	
	loop {
		match vote.next_preference() {
			Some(preference) => {
				let candidate = &state.election.candidates[preference];
				let count_card = &state.candidates[candidate];
				
				if let CandidateState::Hopeful | CandidateState::Guarded = count_card.state {
					next_candidate = Some(candidate);
					break;
				}
			}
			None => { break; }
		}
	}
	
	// Have to structure like this to satisfy Rust's borrow checker
	if let Some(candidate) = next_candidate {
		// Available preference
		state.candidates.get_mut(source_candidate).unwrap().transfer(&-vote.ballot.orig_value.clone());
		
		let count_card = state.candidates.get_mut(candidate).unwrap();
		count_card.transfer(&vote.ballot.orig_value);
		
		match count_card.parcels.last_mut() {
			Some(parcel) => {
				if parcel.source_order == state.num_elected + state.num_excluded {
					parcel.votes.push(vote);
				} else {
					let parcel = Parcel {
						votes: vec![vote],
						value_fraction: N::one(),
						source_order: state.num_elected + state.num_excluded,
					};
					count_card.parcels.push(parcel);
				}
			}
			None => {
				let parcel = Parcel {
					votes: vec![vote],
					value_fraction: N::one(),
					source_order: state.num_elected + state.num_excluded,
				};
				count_card.parcels.push(parcel);
			}
		}
		
		if opts.sample_per_ballot {
			super::elect_hopefuls(state, opts, true)?;
		}
	} else {
		// Exhausted
		if opts.transferable_only && ignore_nontransferable {
			// Another ballot paper required
		} else {
			state.candidates.get_mut(source_candidate).unwrap().transfer(&-vote.ballot.orig_value.clone());
			state.exhausted.transfer(&vote.ballot.orig_value);
			
			match state.exhausted.parcels.last_mut() {
				Some(parcel) => {
					if parcel.source_order == state.num_elected + state.num_excluded {
						parcel.votes.push(vote);
					} else {
						let parcel = Parcel {
							votes: vec![vote],
							value_fraction: N::one(),
							source_order: state.num_elected + state.num_excluded,
						};
						state.exhausted.parcels.push(parcel);
					}
				}
				None => {
					let parcel = Parcel {
						votes: vec![vote],
						value_fraction: N::one(),
						source_order: state.num_elected + state.num_excluded,
					};
					state.exhausted.parcels.push(parcel);
				}
			}
		}
	}
	
	return Ok(());
}

/// Perform one stage of a candidate exclusion according to the random subset method
pub fn exclude_candidates<'a, N: Number>(state: &mut CountState<'a, N>, opts: &STVOptions, excluded_candidates: Vec<&'a Candidate>) -> Result<(), STVError>
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);
		}
	}
	
	// Count votes
	let mut total_ballots: usize = 0;
	for excluded_candidate in excluded_candidates.iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		total_ballots = count_card.parcels.iter()
			.fold(total_ballots, |acc, p| acc + p.votes.len());
	}
	
	if total_ballots == 1 {
		state.logger.log_literal("Transferring 1 ballot.".to_string());
	} else {
		state.logger.log_literal(format!("Transferring {:.0} ballots.", total_ballots));
	}
	
	// Transfer votes
	for excluded_candidate in excluded_candidates.iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		let votes = count_card.concat_parcels();
		
		for vote in votes {
			transfer_ballot(state, opts, excluded_candidate, vote, false)?;
			if state.num_elected == state.election.seats {
				return Ok(());
			}
		}
		
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		count_card.finalised = true;
	}
	
	return Ok(());
}