/*  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, Vote};
use crate::numbers::Number;
use crate::stv::{STVOptions, SurplusMethod};

use std::collections::HashMap;
use std::ops;

use super::STVError;

/// Distribute the surplus of a given candidate according to the random subset method, based on [STVOptions::surplus]
pub fn distribute_surplus<N: Number>(state: &mut CountState<N>, opts: &STVOptions, elected_candidate: &Candidate) -> Result<(), 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>
{
	state.kind = Some("Surplus of");
	state.title = String::from(&elected_candidate.name);
	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 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!()
	}
	
	// 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() {
		let mut vote = numbered_votes.remove(&index).unwrap();
		
		// Transfer to next preference
		let mut next_candidate = None;
		for (i, preference) in vote.ballot.preferences.iter().enumerate().skip(vote.up_to_pref) {
			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);
				vote.up_to_pref = i + 1;
				break;
			}
		}
		
		// Have to structure like this to satisfy Rust's borrow checker
		if let Some(candidate) = next_candidate {
			// Available preference
			state.candidates.get_mut(elected_candidate).unwrap().transfer(&-vote.value.clone());
			
			let count_card = state.candidates.get_mut(candidate).unwrap();
			count_card.transfer(&vote.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],
							source_order: state.num_elected + state.num_excluded,
						};
						count_card.parcels.push(parcel);
					}
				}
				None => {
					let parcel = Parcel {
						votes: vec![vote],
						source_order: state.num_elected + state.num_excluded,
					};
					count_card.parcels.push(parcel);
				}
			}
			
			if opts.sample_per_ballot {
				super::elect_hopefuls(state, opts)?;
			}
		} else {
			// Exhausted
			if opts.transferable_only {
				// Another ballot paper required
			} else {
				state.candidates.get_mut(elected_candidate).unwrap().transfer(&-vote.value.clone());
				state.exhausted.transfer(&vote.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],
								source_order: state.num_elected + state.num_excluded,
							};
							state.exhausted.parcels.push(parcel);
						}
					}
					None => {
						let parcel = Parcel {
							votes: vec![vote],
							source_order: state.num_elected + state.num_excluded,
						};
						state.exhausted.parcels.push(parcel);
					}
				}
			}
		}
		
		index += skip_value;
		if index >= total_ballots {
			iteration += 1;
			index = iteration + skip_value - 1;
			
			if iteration >= skip_value {
				// 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;
			}
		}
	}
	
	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);
		}
	}
	
	// Determine votes to transfer
	let mut votes = Vec::new();
	
	for excluded_candidate in excluded_candidates.iter() {
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		votes.append(&mut count_card.concat_parcels());
		count_card.parcels.clear();
	}
	
	let total_ballots = votes.len();
	
	if !votes.is_empty() {
		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 vote by vote
		for mut vote in votes {
			// Subtract votes from excluded candidate
			let count_card = state.candidates.get_mut(&state.election.candidates[vote.ballot.preferences[vote.up_to_pref - 1]]).unwrap();
			count_card.transfer(&-vote.value.clone());
			
			// Transfer to next preference
			let mut next_candidate = None;
			for (i, preference) in vote.ballot.preferences.iter().enumerate().skip(vote.up_to_pref) {
				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);
					vote.up_to_pref = i + 1;
					break;
				}
			}
			
			// Have to structure like this to satisfy Rust's borrow checker
			if let Some(candidate) = next_candidate {
				// Available preference
				let count_card = state.candidates.get_mut(candidate).unwrap();
				count_card.transfer(&vote.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],
								source_order: state.num_elected + state.num_excluded,
							};
							count_card.parcels.push(parcel);
						}
					}
					None => {
						let parcel = Parcel {
							votes: vec![vote],
							source_order: state.num_elected + state.num_excluded,
						};
						count_card.parcels.push(parcel);
					}
				}
				
				if opts.sample_per_ballot {
					super::elect_hopefuls(state, opts)?;
				}
			} else {
				// Exhausted
				state.exhausted.transfer(&vote.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],
								source_order: state.num_elected + state.num_excluded,
							};
							state.exhausted.parcels.push(parcel);
						}
					}
					None => {
						let parcel = Parcel {
							votes: vec![vote],
							source_order: state.num_elected + state.num_excluded,
						};
						state.exhausted.parcels.push(parcel);
					}
				}
			}
		}
	}
	
	return Ok(());
}