/*  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/>.
 */

#![allow(mutable_borrow_reservation_conflict)]

//#[cfg(target_arch = "wasm32")]
pub mod wasm;

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

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

pub struct STVOptions<'a> {
	pub round_votes: Option<usize>,
	pub exclusion: &'a str,
}

pub fn count_init<N: Number>(mut state: &mut CountState<'_, N>, _opts: &STVOptions) {
	distribute_first_preferences(&mut state);
	calculate_quota(&mut state);
	elect_meeting_quota(&mut state);
}

pub fn count_one_stage<N: Number>(mut state: &mut CountState<'_, N>, opts: &STVOptions) -> bool
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.entries.clear();
	state.step_all();
	
	// Finish count
	if finished_before_stage(&state) {
		return true;
	}
	
	// Continue exclusions
	if continue_exclusion(&mut state, &opts) {
		elect_meeting_quota(&mut state);
		return false;
	}
	
	// Distribute surpluses
	if distribute_surpluses(&mut state, &opts) {
		elect_meeting_quota(&mut state);
		return false;
	}
	
	// Attempt bulk election
	if bulk_elect(&mut state) {
		elect_meeting_quota(&mut state);
		return false;
	}
	
	// Exclude lowest hopeful
	if exclude_hopefuls(&mut state, &opts) {
		elect_meeting_quota(&mut state);
		return false;
	}
	
	todo!();
}

struct NextPreferencesResult<'a, N> {
	candidates: HashMap<&'a Candidate, NextPreferencesEntry<'a, N>>,
	exhausted: NextPreferencesEntry<'a, N>,
	total_ballots: N,
	total_votes: N,
}

struct NextPreferencesEntry<'a, N> {
	//count_card: Option<&'a CountCard<'a, N>>,
	votes: Vec<Vote<'a, N>>,
	num_ballots: N,
	num_votes: N,
}

fn next_preferences<'a, N: Number>(state: &CountState<'a, N>, votes: Vec<Vote<'a, N>>) -> NextPreferencesResult<'a, N> {
	let mut result = NextPreferencesResult {
		candidates: HashMap::new(),
		exhausted: NextPreferencesEntry {
			votes: Vec::new(),
			num_ballots: N::new(),
			num_votes: N::new(),
		},
		total_ballots: N::new(),
		total_votes: N::new(),
	};
	
	for mut vote in votes.into_iter() {
		result.total_ballots += &vote.ballot.orig_value;
		result.total_votes += &vote.value;
		
		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.get(candidate).unwrap();
			
			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 {
			if result.candidates.contains_key(candidate) {
				let entry = result.candidates.get_mut(candidate).unwrap();
				entry.num_ballots += &vote.ballot.orig_value;
				entry.num_votes += &vote.value;
				entry.votes.push(vote);
			} else {
				let entry = NextPreferencesEntry {
					num_ballots: vote.ballot.orig_value.clone(),
					num_votes: vote.value.clone(),
					votes: vec![vote],
				};
				result.candidates.insert(candidate, entry);
			}
		} else {
			result.exhausted.num_ballots += &vote.ballot.orig_value;
			result.exhausted.num_votes += &vote.value;
			result.exhausted.votes.push(vote);
		}
	}
	
	return result;
}

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 = 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());
}

fn calculate_quota<N: Number>(state: &mut CountState<N>) {
	let mut log = String::new();
	
	// Calculate the total vote
	state.quota = state.candidates.values().fold(N::zero(), |acc, cc| { acc + &cc.votes });
	log.push_str(format!("{:.2} usable votes, so the quota is ", state.quota).as_str());
	
	// TODO: Different quotas
	state.quota /= N::from(state.election.seats + 1);
	
	// TODO: Different rounding rules
	state.quota += N::one();
	state.quota.floor_mut(0);
	log.push_str(format!("{:.2}.", state.quota).as_str());
	
	state.logger.log_literal(log);
}

fn meets_quota<N: Number>(quota: &N, count_card: &CountCard<N>) -> bool {
	// TODO: Different quota rules
	return count_card.votes >= *quota;
}

fn elect_meeting_quota<N: Number>(state: &mut CountState<N>) {
	let quota = &state.quota; // Have to do this or else the borrow checker gets confused
	let mut cands_meeting_quota: Vec<(&&Candidate, &mut CountCard<N>)> = state.candidates.iter_mut()
		.filter(|(_, cc)| cc.state == CandidateState::HOPEFUL && meets_quota(quota, cc))
		.collect();
	
	if cands_meeting_quota.len() > 0 {
		// Sort by votes
		cands_meeting_quota.sort_unstable_by(|a, b| a.1.votes.partial_cmp(&b.1.votes).unwrap());
		
		// Declare elected in descending order of votes
		for (candidate, count_card) in cands_meeting_quota.into_iter().rev() {
			count_card.state = CandidateState::ELECTED;
			state.num_elected += 1;
			count_card.order_elected = state.num_elected as isize;
			state.logger.log_smart(
				"{} meets the quota and is elected.",
				"{} meet the quota and are elected.",
				vec![&candidate.name]
			);
		}
	}
}

fn distribute_surpluses<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> bool
where
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	let mut has_surplus: Vec<(&&Candidate, &CountCard<N>)> = state.candidates.iter()
		.filter(|(_, cc)| cc.votes > state.quota)
		.collect();
	
	if has_surplus.len() > 0 {
		// TODO: Different sorting orders
		has_surplus.sort_unstable_by(|a, b| a.1.order_elected.partial_cmp(&b.1.order_elected).unwrap());
		
		// Distribute top candidate's surplus
		// TODO: Handle ties
		let elected_candidate = has_surplus.first_mut().unwrap().0;
		distribute_surplus(state, &opts, elected_candidate);
		
		return true;
	}
	return false;
}

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::Neg<Output=N>
{
	let count_card = state.candidates.get(elected_candidate).unwrap();
	let surplus = &count_card.votes - &state.quota;
	
	// Inclusive Gregory
	// TODO: Other methods
	let votes = state.candidates.get(elected_candidate).unwrap().parcels.concat();
	
	// Count next preferences
	let result = next_preferences(state, votes);
	
	// Transfer candidate votes
	// Unweighted inclusive Gregory
	// TODO: Other methods
	let transfer_value = surplus.clone() / &result.total_ballots;
	
	state.kind = Some("Surplus of");
	state.title = String::from(&elected_candidate.name);
	state.logger.log_literal(format!("Surplus of {} distributed at value {:.2}.", elected_candidate.name, transfer_value));
	
	let mut checksum = N::new();
	
	for (candidate, entry) in result.candidates.into_iter() {
		let mut parcel = entry.votes as Parcel<N>;
		
		// Reweight votes
		for vote in parcel.iter_mut() {
			//vote.value = vote.ballot.orig_value.clone() * &transfer_value;
			vote.value = vote.ballot.orig_value.clone() * &surplus / &result.total_ballots;
		}
		
		let count_card = state.candidates.get_mut(candidate).unwrap();
		count_card.parcels.push(parcel);
		
		let mut candidate_transfers = entry.num_ballots * &surplus / &result.total_ballots;
		// Round transfers
		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_ballots * &surplus / &result.total_ballots;
	if let Some(dps) = opts.round_votes {
		exhausted_transfers.floor_mut(dps);
	}
	state.exhausted.transfer(&exhausted_transfers);
	checksum += exhausted_transfers;
	
	// Finalise candidate votes
	let count_card = state.candidates.get_mut(elected_candidate).unwrap();
	count_card.transfers = -&surplus;
	count_card.votes.assign(&state.quota);
	checksum -= surplus;
	
	// Update loss by fraction
	state.loss_fraction.transfer(&-checksum);
}

fn bulk_elect<N: Number>(state: &mut CountState<N>) -> bool {
	if state.election.candidates.len() - state.num_excluded <= state.election.seats {
		state.kind = None;
		state.title = "Bulk election".to_string();
		
		// Bulk elect all remaining candidates
		let mut hopefuls: Vec<(&&Candidate, &mut CountCard<N>)> = state.candidates.iter_mut()
			.filter(|(_, cc)| cc.state == CandidateState::HOPEFUL)
			.collect();
		
		// TODO: Handle ties
		hopefuls.sort_unstable_by(|a, b| a.1.votes.partial_cmp(&b.1.votes).unwrap());
		
		for (candidate, count_card) in hopefuls.into_iter() {
			count_card.state = CandidateState::ELECTED;
			state.num_elected += 1;
			count_card.order_elected = state.num_elected as isize;
			
			state.logger.log_smart(
				"{} is elected to fill the remaining vacancy.",
				"{} are elected to fill the remaining vacancies.",
				vec![&candidate.name]
			);
		}
		
		return true;
	}
	return false;
}

fn exclude_hopefuls<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> bool
where
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
{
	let mut hopefuls: Vec<(&&Candidate, &CountCard<N>)> = state.candidates.iter()
		.filter(|(_, cc)| cc.state == CandidateState::HOPEFUL)
		.collect();
	
	// Sort by votes
	// TODO: Handle ties
	hopefuls.sort_unstable_by(|a, b| a.1.votes.partial_cmp(&b.1.votes).unwrap());
	
	// Exclude lowest ranked candidate
	let excluded_candidate = hopefuls.first().unwrap().0;
	
	state.kind = Some("Exclusion of");
	state.title = String::from(&excluded_candidate.name);
	state.logger.log_smart(
		"No surpluses to distribute, so {} is excluded.",
		"No surpluses to distribute, so {} are excluded.",
		vec![&excluded_candidate.name]
	);
	
	exclude_candidate(state, opts, excluded_candidate);
	
	return true;
}

fn continue_exclusion<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> bool
where
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
{
	// Cannot filter by raw vote count, as candidates may have 0.00 votes but still have recorded ballot papers
	let mut excluded_with_votes: Vec<(&&Candidate, &CountCard<N>)> = state.candidates.iter()
		//.filter(|(_, cc)| cc.state == CandidateState::EXCLUDED && !cc.votes.is_zero())
		.filter(|(_, cc)| cc.state == CandidateState::EXCLUDED && cc.parcels.iter().any(|p| p.len() > 0))
		.collect();
	
	if excluded_with_votes.len() > 0 {
		excluded_with_votes.sort_unstable_by(|a, b| a.1.order_elected.partial_cmp(&b.1.order_elected).unwrap());
		let excluded_candidate = excluded_with_votes.first().unwrap().0;
		
		state.kind = Some("Exclusion of");
		state.title = String::from(&excluded_candidate.name);
		state.logger.log_smart(
			"Continuing exclusion of {}.",
			"Continuing exclusion of {}.",
			vec![&excluded_candidate.name]
		);
		
		exclude_candidate(state, opts, excluded_candidate);
		return true;
	}
	
	return false;
}

fn exclude_candidate<N: Number>(state: &mut CountState<N>, opts: &STVOptions, excluded_candidate: &Candidate)
where
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
{
	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 = -(state.num_excluded as isize);
	}
	
	// Determine votes to transfer in this stage
	let mut votes;
	let votes_remaining;
	
	if opts.exclusion == "one_round" {
		// Exclude in one round
		votes = count_card.parcels.concat();
		votes_remaining = 0;
		
	} else if opts.exclusion == "by_value" {
		// Exclude by value
		let all_votes = count_card.parcels.concat();
		
		// TODO: Write a multiple min/max function
		let min_value = all_votes.iter().map(|v| &v.value / &v.ballot.orig_value).max().unwrap();
		
		votes = Vec::new();
		let mut remaining_votes = Vec::new();
		
		// This could be implemented using Vec.drain_filter, but that is experimental currently
		for vote in all_votes.into_iter() {
			if &vote.value / &vote.ballot.orig_value == min_value {
				votes.push(vote);
			} else {
				remaining_votes.push(vote);
			}
		}
		
		votes_remaining = remaining_votes.len();
		// Leave remaining votes with candidate (as one parcel)
		count_card.parcels = vec![remaining_votes];
		
	} else {
		// TODO: Exclude by parcel
		panic!("Invalid --exclusion");
	}
	
	let mut checksum = N::new();
	
	if votes.len() > 0 {
		let value = &votes[0].value / &votes[0].ballot.orig_value;
		
		// Count next preferences
		let result = next_preferences(state, votes);
		
		if opts.exclusion == "one_round" {
			state.logger.log_literal(format!("Transferring {:.0} ballot papers, totalling {:.2} votes.", result.total_ballots, result.total_votes));
		} else if opts.exclusion == "by_value" {
			state.logger.log_literal(format!("Transferring {:.0} ballot papers, totalling {:.2} votes, received at value {:.2}.", result.total_ballots, result.total_votes, value));
		}
		
		// 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_remaining > 0 {
			// Subtract from candidate tally
			let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
			checksum -= &result.total_votes;
			count_card.transfer(&-result.total_votes);
			
			// By definition, there is no loss by fraction
		}
	}
	
	if votes_remaining == 0 {
		// Finalise candidate votes
		let count_card = state.candidates.get_mut(excluded_candidate).unwrap();
		checksum -= &count_card.votes;
		count_card.transfers = -count_card.votes.clone();
		count_card.votes = N::new();
		
		// Update loss by fraction
		state.loss_fraction.transfer(&-checksum);
		
		if opts.exclusion != "one_round" {
			state.logger.log_literal("Exclusion complete.".to_string());
		}
	}
}

fn finished_before_stage<N: Number>(state: &CountState<N>) -> bool {
	if state.num_elected >= state.election.seats {
		return true;
	}
	return false;
}