/*  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::{Constraints, ConstraintMatrix};
use crate::logger::Logger;
use crate::numbers::Number;
use crate::sharandom::SHARandom;

use std::collections::HashMap;

/// An election to be counted
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>>,
	/// Constraints on candidates
	pub constraints: Option<Constraints>,
}

impl<N: Number> Election<N> {
	/// Parse the given BLT file and return an [Election]
	pub fn from_blt<I: Iterator<Item=String>>(mut lines: I) -> Self {
		// Read first line
		let line = lines.next().expect("Unexpected EOF");
		let mut bits = line.split(" ");
		let num_candidates = bits.next().expect("Syntax Error").parse().expect("Syntax Error");
		let seats: usize = bits.next().expect("Syntax Error").parse().expect("Syntax Error");
		
		// Initialise the Election object
		let mut election = Election {
			name: String::new(),
			seats: seats,
			candidates: Vec::with_capacity(num_candidates),
			withdrawn_candidates: Vec::new(),
			ballots: Vec::new(),
			constraints: None,
		};
		
		// Read ballots
		for line in &mut lines {
			if line == "0" {
				break;
			}
			
			let mut bits = line.split(" ");
			
			if line.starts_with("-") {
				// Withdrawn candidates
				for bit in bits.into_iter() {
					let val = bit[1..bit.len()].parse::<usize>().expect("Syntax Error");
					election.withdrawn_candidates.push(val - 1);
				}
				continue;
			}
			
			let value = N::parse(bits.next().expect("Syntax Error"));
			
			let mut ballot = Ballot {
				orig_value: value,
				preferences: Vec::new(),
			};
			
			for preference in bits {
				if preference != "0" {
					let preference = preference.parse::<usize>().expect("Syntax Error");
					ballot.preferences.push(preference - 1);
				}
			}
			
			election.ballots.push(ballot);
		}
		
		// Read candidates
		for line in lines.by_ref().take(num_candidates) {
			let mut line = &line[..];
			if line.starts_with("\"") && line.ends_with("\"") {
				line = &line[1..line.len()-1];
			}
			
			election.candidates.push(Candidate { name: line.to_string() });
		}
		
		// Read name
		let line = lines.next().expect("Syntax Error");
		let mut line = &line[..];
		if line.starts_with("\"") && line.ends_with("\"") {
			line = &line[1..line.len()-1];
		}
		election.name.push_str(line);
		
		return election;
	}
	
	/// Convert ballots with weight >1 to multiple ballots of weight 1
	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(),
				};
				normalised_ballots.push(new_ballot);
				n += &one;
			}
		}
		self.ballots = normalised_ballots;
	}
}

/// A candidate in an [Election]
#[derive(PartialEq, Eq, Hash)]
pub struct Candidate {
	/// Name of the candidate
	pub name: String,
}

/// 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>,
	
	/// [HashMap] of [CountCard]s for each [Candidate] in the election
	pub candidates: HashMap<&'a Candidate, CountCard<'a, N>>,
	/// [CountCard] representing the exhausted pile
	pub exhausted: CountCard<'a, N>,
	/// [CountCard] representing loss by fraction
	pub loss_fraction: CountCard<'a, N>,
	
	/// [crate::stv::meek::BallotTree] for Meek STV
	pub ballot_tree: Option<crate::stv::meek::BallotTree<'a, N>>,
	
	/// Values used to break ties, based on forwards tie-breaking
	pub forwards_tiebreak: Option<HashMap<&'a Candidate, usize>>,
	/// Values used to break ties, based on backwards tie-breaking
	pub backwards_tiebreak: Option<HashMap<&'a Candidate, usize>>,
	/// [SHARandom] for random tie-breaking
	pub random: Option<SHARandom<'a>>,
	
	/// Quota for election
	pub quota: Option<N>,
	/// Vote required for election
	///
	/// With a static quota, this is equal to the quota. With ERS97 rules, this may vary from the quota.
	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>,
	
	/// The type of stage being counted
	///
	/// For example, "Surplus of", "Exclusion of"
	pub kind: Option<&'a str>,
	/// The description of the stage being counted, excluding [CountState::kind]
	pub title: String,
	/// [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: &election,
			candidates: HashMap::new(),
			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,
			kind: None,
			title: String::new(),
			logger: Logger { entries: Vec::new() },
		};
		
		for candidate in election.candidates.iter() {
			state.candidates.insert(candidate, CountCard::new());
		}
		
		for withdrawn_idx in election.withdrawn_candidates.iter() {
			state.candidates.get_mut(&election.candidates[*withdrawn_idx]).unwrap().state = CandidateState::Withdrawn;
		}
		
		if let Some(constraints) = &election.constraints {
			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();
	}
}

/// Represents either a reference to a [CountState] or a clone
#[allow(dead_code)]
pub enum CountStateOrRef<'a, N: Number> {
	/// Cloned [CountState]
	///
	/// Currently unused/unimplemented, but may be used in future for rollback-based constraints
	State(CountState<'a, N>),
	/// Reference to a [CountState]
	Ref(&'a CountState<'a, N>),
}

impl<'a, N: Number> CountStateOrRef<'a, N> {
	/// Construct a [CountStateOrRef] as a reference to a [CountState]
	pub fn from(state: &'a CountState<N>) -> Self {
		return Self::Ref(state);
	}
	
	/// Return a reference to the underlying [CountState]
	pub fn as_ref(&self) -> &CountState<N> {
		match self {
			CountStateOrRef::State(state) => &state,
			CountStateOrRef::Ref(state) => state,
		}
	}
}

/// Result of a stage of counting
pub struct StageResult<'a, N: Number> {
	/// See [CountState::kind]
	pub kind: Option<&'a str>,
	/// See [CountState::title]
	pub title: &'a String,
	/// Detailed logs of this stage, rendered from [CountState::logger]
	pub logs: Vec<String>,
	/// Reference to the [CountState] or cloned [CountState] of this stage
	pub state: CountStateOrRef<'a, N>,
}

/// 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,
	
	//pub orig_votes: N,
	/// 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,
	
	/// 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,
			//orig_votes: N::new(),
			transfers: N::new(),
			votes: 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.orig_votes = self.votes.clone();
		self.transfers = N::new();
	}
}

/// Parcel of [Vote]s during a count
pub type Parcel<'a, N> = Vec<Vote<'a, N>>;

/// 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>,
	/// Current value of the ballot
	pub value: N,
	/// Index of the next preference to examine
	pub up_to_pref: usize,
}

/// A record of a voter's preferences
pub struct Ballot<N> {
	/// Original value/weight of the ballot
	pub orig_value: N,
	/// Indexes of candidates preferenced on the ballot
	pub preferences: Vec<usize>,
}

/// State of a [Candidate] during a count
#[allow(dead_code)]
#[derive(PartialEq)]
#[derive(Clone)]
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,
}