/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021–2022 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::candmap::CandidateMap;
use crate::election::{Candidate, CandidateState, CountCard, CountState, Election, StageKind, RollbackState};
use crate::numbers::Number;
use crate::stv::{self, gregory, sample, ConstraintMode, STVError, STVOptions, SurplusMethod, SurplusOrder};
use crate::ties::{self, TieStrategy};

use itertools::Itertools;
use ndarray::{Array, Dimension, IxDyn};

#[cfg(not(target_arch = "wasm32"))]
use rkyv::{Archive, Deserialize, Serialize};

use std::fmt;
use std::num::ParseIntError;
use std::ops;

/// Constraints for an [crate::election::Election]
#[derive(Clone, Debug)]
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
pub struct Constraints(pub Vec<Constraint>);

impl Constraints {
	/// Parse the given CON file and return a [Constraints]
	pub fn from_con<S: AsRef<str>, I: Iterator<Item=S>>(lines: I) -> Result<Self, ParseError> {
		let mut constraints = Constraints(Vec::new());
		
		for (line_no, line) in lines.enumerate() {
			let mut bits = line.as_ref().split(' ').peekable();
			
			// Read constraint category and group
			let constraint_name = read_quoted_string(line_no, &mut bits)?;
			let group_name = read_quoted_string(line_no, &mut bits)?;
			
			// Read min, max
			let min: usize = bits
				.next().ok_or(ParseError::UnexpectedEOL(line_no, "minimum number"))?
				.parse().map_err(|e| ParseError::InvalidNumber(line_no, e))?;
			let max: usize = bits
				.next().ok_or(ParseError::UnexpectedEOL(line_no, "maximum number"))?
				.parse().map_err(|e| ParseError::InvalidNumber(line_no, e))?;
			
			// Read candidates
			let mut candidates: Vec<usize> = Vec::new();
			for x in bits {
				candidates.push(x.parse::<usize>().map_err(|e| ParseError::InvalidNumber(line_no, e))? - 1);
			}
			
			// Insert constraint/group
			let constraint = match constraints.0.iter_mut().find(|c| c.name == constraint_name) {
				Some(c) => { c }
				None => {
					let c = Constraint {
						name: constraint_name,
						groups: Vec::new(),
					};
					constraints.0.push(c);
					constraints.0.last_mut().unwrap()
				}
			};
			
			if constraint.groups.iter().any(|g| g.name == group_name) {
				return Err(ParseError::DuplicateGroup(line_no, group_name, constraint.name.clone()));
			}
			
			constraint.groups.push(ConstrainedGroup {
				name: group_name,
				candidates,
				min,
				max,
			});
		}
		
		return Ok(constraints);
	}
	
	/// Validate that each candidate is specified exactly once in each constraint, and (if applicable) limitations of the constraint mode are applied
	pub fn validate_constraints(&self, num_candidates: usize, constraint_mode: ConstraintMode) -> Result<(), ValidationError> {
		for constraint in &self.0 {
			let mut remaining_candidates: Vec<usize> = (0..num_candidates).collect();
			
			for group in &constraint.groups {
				for candidate in &group.candidates {
					match remaining_candidates.iter().position(|c| c == candidate) {
						Some(idx) => {
							remaining_candidates.remove(idx);
						}
						None => {
							return Err(ValidationError::DuplicateCandidate(*candidate, constraint.name.clone()));
						}
					}
				}
				
				if constraint_mode == ConstraintMode::RepeatCount {
					// Each group must be either a maximum constraint, or the remaining group
					if group.min == 0 {
						// Maximum constraint: OK
					} else if group.max >= group.candidates.len() {
						// Remaining group: OK
					} else {
						return Err(ValidationError::InvalidTwoStage(constraint.name.clone(), group.name.clone()));
					}
					
					// FIXME: Is other validation required?
				}
			}
			
			if !remaining_candidates.is_empty() {
				return Err(ValidationError::UnassignedCandidate(*remaining_candidates.first().unwrap(), constraint.name.clone()));
			}
		}
		
		Ok(())
	}
	
	/// Check if any elected candidates exceed constrained maximums
	pub fn exceeds_maximum<'a, N: Number>(&self, election: &Election<N>, candidates: CandidateMap<CountCard<'a, N>>) -> Option<(&Constraint, &ConstrainedGroup)> {
		for constraint in &self.0 {
			for group in &constraint.groups {
				let mut num_elected = 0;
				for candidate in &group.candidates {
					if candidates[&election.candidates[*candidate]].state == CandidateState::Elected {
						num_elected += 1;
					}
				}
				if num_elected > group.max {
					return Some((&constraint, &group));
				}
			}
		}
		return None;
	}
}

/// Error parsing constraints
pub enum ParseError {
	/// Duplicate group in a constraint
	DuplicateGroup(usize, String, String),
	/// Unexpected EOL, expected ...
	UnexpectedEOL(usize, &'static str),
	/// Invalid number
	InvalidNumber(usize, ParseIntError),
}

impl fmt::Display for ParseError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		match self {
			ParseError::DuplicateGroup(line_no, group_name, constraint_name) => {
				f.write_fmt(format_args!(r#"Line {}, duplicate group "{}" in constraint "{}""#, line_no, group_name, constraint_name))
			}
			ParseError::UnexpectedEOL(line_no, expected) => {
				f.write_fmt(format_args!(r#"Line {}, unexpected end-of-line, expected {}"#, line_no, expected))
			}
			ParseError::InvalidNumber(line_no, err) => {
				f.write_fmt(format_args!(r#"Line {}, invalid number: {}"#, line_no, err))
			}
		}
	}
}

impl fmt::Debug for ParseError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		return fmt::Display::fmt(self, f);
	}
}

/// Error validating constraints
pub enum ValidationError {
	/// Duplicate candidate in a constraint
	DuplicateCandidate(usize, String),
	/// Unassigned candidate in a constraint
	UnassignedCandidate(usize, String),
	/// Constraint is incompatible with ConstraintMode::TwoStage
	InvalidTwoStage(String, String),
}

impl fmt::Display for ValidationError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		match self {
			ValidationError::DuplicateCandidate(candidate, constraint_name) => {
				f.write_fmt(format_args!(r#"Duplicate candidate {} in constraint "{}""#, candidate + 1, constraint_name))
			}
			ValidationError::UnassignedCandidate(candidate, constraint_name) => {
				f.write_fmt(format_args!(r#"Unassigned candidate {} in constraint "{}""#, candidate + 1, constraint_name))
			}
			ValidationError::InvalidTwoStage(constraint_name, group_name) => {
				f.write_fmt(format_args!(r#"Constraint "{}" group "{}" is incompatible with --constraint-mode repeat_count"#, constraint_name, group_name))
			}
		}
	}
}

impl fmt::Debug for ValidationError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		return fmt::Display::fmt(self, f);
	}
}

#[test]
fn duplicate_contraint_group() {
	let input = r#""Constraint 1" "Group 1" 0 3 1 2 3
"Constraint 1" "Group 1" 0 3 4 5 6"#;
	Constraints::from_con(input.lines()).unwrap_err();
}

#[test]
fn duplicate_candidate() {
	let input = r#""Constraint 1" "Group 1" 0 3 1 2 3 4
"Constraint 1" "Group 2" 0 3 4 5 6"#;
	let constraints = Constraints::from_con(input.lines()).unwrap();
	constraints.validate_constraints(6, ConstraintMode::GuardDoom).unwrap_err();
}

#[test]
fn unassigned_candidate() {
	let input = r#""Constraint 1" "Group 1" 0 3 1 2 3
"Constraint 1" "Group 2" 0 3 4 5 6"#;
	let constraints = Constraints::from_con(input.lines()).unwrap();
	constraints.validate_constraints(7, ConstraintMode::GuardDoom).unwrap_err();
}

/// Read an optionally quoted string, returning the string without quotes
fn read_quoted_string<'a, I: Iterator<Item=&'a str>>(line_no: usize, bits: &mut I) -> Result<String, ParseError> {
	let x = bits.next().ok_or(ParseError::UnexpectedEOL(line_no, "string continuation"))?;
	if let Some(x1) = x.strip_prefix('"') {
		if let Some(x2) = x.strip_suffix('"') {
			// Complete string
			return Ok(String::from(x2));
		} else {
			// Incomplete string
			let mut result = String::from(x1);
			
			// Read until matching "
			loop {
				let x = bits.next().ok_or(ParseError::UnexpectedEOL(line_no, "string continuation"))?;
				result.push(' ');
				if let Some(x1) = x.strip_suffix('"') {
					// End of string
					result.push_str(x1);
					break;
				} else {
					// Middle of string
					result.push_str(x);
				}
			}
			
			return Ok(result);
		}
	} else {
		// Unquoted string
		return Ok(String::from(x));
	}
}

/// A single dimension of constraint
#[derive(Clone, Debug)]
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
pub struct Constraint {
	/// Name of this constraint
	pub name: String,
	/// Groups of candidates within this constraint
	pub groups: Vec<ConstrainedGroup>,
}

/// A group of candidates, of which a certain minimum and maximum must be elected
#[derive(Clone, Debug)]
#[cfg_attr(not(target_arch = "wasm32"), derive(Archive, Deserialize, Serialize))]
pub struct ConstrainedGroup {
	/// Name of this group
	pub name: String,
	/// Indexes of [crate::election::Candidate]s to constrain
	pub candidates: Vec<usize>,
	/// Minimum number to elect
	pub min: usize,
	/// Maximum number to elect
	pub max: usize,
}

/// Error reaching a stable state when processing constraints
#[derive(Debug)]
pub enum ConstraintError {
	/// No conformant result is possible
	NoConformantResult,
}

// ----------------------
// GUARD/DOOM CONSTRAINTS
// ----------------------

/// Cell in a [ConstraintMatrix]
#[derive(Clone)]
pub struct ConstraintMatrixCell {
	/// Number of elected candidates in this cell
	pub elected: usize,
	/// Minimum number of candidates which must be elected from this cell for a conformant result
	pub min: usize,
	/// Maximum number of candidates which may be elected from this cell for a conformant result
	pub max: usize,
	/// Total number of elected or hopeful candidates in this cell
	pub cands: usize,
}

/// N-dimensional cube of [ConstraintMatrixCell]s representing the conformant combinations of elected candidates
#[derive(Clone)]
pub struct ConstraintMatrix(pub Array<ConstraintMatrixCell, IxDyn>);

impl ConstraintMatrix {
	/// Return a new [ConstraintMatrix], with the specified number of groups for each constraint dimension
	pub fn new(constraints: &mut [usize]) -> Self {
		// Add 1 to dimensions for totals cells
		for c in constraints.iter_mut() {
			*c += 1;
		}
		
		return Self(Array::from_elem(
			IxDyn(constraints),
			ConstraintMatrixCell {
				elected: 0,
				min: 0,
				max: 0,
				cands: 0,
			}
		));
	}
	
	/// Initialise the matrix once the number of candidates in each innermost cell, and min/max for each constraint group, have been specified
	pub fn init(&mut self) {
		// Compute candidate totals
		self.recount_cands();
		
		// Initialise max for grand total cell
		let idx = IxDyn(&vec![0; self.0.ndim()][..]);
		self.0[&idx].max = self.0[&idx].cands;
		
		// Initialise max for inner cells (>=2 zeroes)
		for idx in ndarray::indices(self.0.shape()) {
			if (0..idx.ndim()).fold(0, |acc, d| if idx[d] != 0 { acc + 1 } else { acc }) < 2 {
				continue;
			}
			self.0[&idx].max = self.0[&idx].cands;
		}
		
		// NB: Bounds on min, max, etc. will be further refined in initial step() calls
	}
	
	/// Update cands/elected in innermost cells based on the provided [CountState::candidates](crate::election::CountState::candidates)
	pub fn update_from_state<N: Number>(&mut self, election: &Election<N>, candidates: &CandidateMap<CountCard<N>>) {
		let constraints = election.constraints.as_ref().unwrap();
		
		// Reset innermost cells
		for idx in ndarray::indices(self.0.shape()) {
			if (0..idx.ndim()).fold(0, |acc, d| if idx[d] == 0 { acc + 1 } else { acc }) != 0 {
				continue;
			}
			self.0[&idx].cands = 0;
			self.0[&idx].elected = 0;
		}
		
		for (i, candidate) in election.candidates.iter().enumerate() {
			if candidate.is_dummy {
				continue;
			}
			
			let idx: Vec<usize> = constraints.0.iter().map(|c| {
				for (j, group) in c.groups.iter().enumerate() {
					if group.candidates.contains(&i) {
						return j + 1;
					}
				}
				// Should be caught by validate_constraints
				unreachable!("Candidate \"{}\" not represented in constraint \"{}\"", candidate.name, c.name);
			}).collect();
			let cell = &mut self[&idx[..]];
			
			let count_card = &candidates[candidate];
			match count_card.state {
				CandidateState::Hopeful | CandidateState::Guarded => {
					cell.cands += 1;
				}
				CandidateState::Elected => {
					cell.cands += 1;
					cell.elected += 1;
				}
				CandidateState::Withdrawn | CandidateState::Doomed | CandidateState::Excluded => {}
			}
		}
	}
	
	/// Recompute [cands](ConstraintMatrixCell::cands) and [elected](ConstraintMatrixCell::elected) for totals cells based on the innermost cells
	pub fn recount_cands(&mut self) {
		let shape = Vec::from(self.0.shape());
		
		// Compute cands/elected totals
		for nzeroes in 1..self.0.ndim()+1 {
			for idx in ndarray::indices(self.0.shape()) {
				// First compute totals cells with 1 zero, then 2 zeroes, ... then grand total (all zeroes)
				if (0..idx.ndim()).fold(0, |acc, d| if idx[d] == 0 { acc + 1 } else { acc }) != nzeroes {
					continue;
				}
				
				self.0[&idx].cands = 0;
				self.0[&idx].elected = 0;
				
				// The axis along which to sum - if multiple, just pick the first, as these should agree
				let zero_axis = (0..idx.ndim()).find(|d| idx[*d] == 0).unwrap();
				
				// Traverse along the axis and sum the candidates
				let mut idx2 = idx.clone();
				for coord in 1..shape[zero_axis] {
					idx2[zero_axis] = coord;
					self.0[&idx].cands += self.0[&idx2].cands;
					self.0[&idx].elected += self.0[&idx2].elected;
				}
			}
		}
	}
	
	/// Attempt to advance the matrix one step towards a stable state
	///
	/// Returns `Ok(true)` if in a stable state, `Ok(false)` if not, and `ConstraintError` on an error.
	///
	pub fn step(&mut self) -> Result<bool, ConstraintError> {
		let shape = Vec::from(self.0.shape());
		
		for idx in ndarray::indices(self.0.shape()) {
			let cell = &mut self.0[&idx];
			
			// Rule 1: Ensure elected < min < max < cands
			if cell.min < cell.elected {
				cell.min = cell.elected;
				return Ok(false);
			}
			if cell.max > cell.cands {
				cell.max = cell.cands;
				return Ok(false);
			}
			if cell.min > cell.max {
				return Err(ConstraintError::NoConformantResult);
			}
			
			let nzeroes = (0..idx.ndim()).fold(0, |acc, d| if idx[d] == 0 { acc + 1 } else { acc });
			
			// Rule 2/3: Ensure min/max is possible in innermost cells
			if nzeroes == 0 {
				for axis in 0..self.0.ndim() {
					let mut idx2 = idx.clone();
					
					// What is the min/max number of candidates that can be elected from other cells in this axis?
					let (other_max, other_min) = (1..shape[axis]).fold((0, 0), |(acc_max, acc_min), coord| {
						if coord == idx[axis] {
							return (acc_max, acc_min);
						}
						idx2[axis] = coord;
						return (acc_max + self.0[&idx2].max, acc_min + self.0[&idx2].min);
					});
					
					// What is the min/max number of candidates that can be elected along this axis?
					idx2[axis] = 0;
					let axis_max = self.0[&idx2].max;
					let axis_min = self.0[&idx2].min;
					
					// How many must be elected from this cell?
					let this_max = (axis_max as i32) - (other_min as i32);
					let this_min = (axis_min as i32) - (other_max as i32);
					
					if this_max < (self.0[&idx].max as i32) {
						self.0[&idx].max = this_max as usize;
						return Ok(false);
					}
					if this_min > (self.0[&idx].min as i32) {
						self.0[&idx].min = this_min as usize;
						return Ok(false);
					}
				}
			}
			
			// Rule 4/5: Ensure min/max is possible in totals cells
			if nzeroes > 0 {
				for axis in 0..self.0.ndim() {
					if idx[axis] != 0 {
						continue;
					}
					
					// What is the total min/max along this axis?
					let mut idx2 = idx.clone();
					let (axis_max, axis_min) = (1..shape[axis]).fold((0, 0), |(acc_max, acc_min), coord| {
						idx2[axis] = coord;
						return (acc_max + self.0[&idx2].max, acc_min + self.0[&idx2].min);
					});
					
					if axis_max < self.0[&idx].max {
						self.0[&idx].max = axis_max;
						return Ok(false);
					}
					if axis_min > self.0[&idx].min {
						self.0[&idx].min = axis_min;
						return Ok(false);
					}
				}
			}
		}
		
		return Ok(true);
	}
}

impl fmt::Display for ConstraintMatrix {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
		let shape = self.0.shape();
		
		let mut result = String::new();
		
		// TODO: >2 dimensions
		if shape.len() == 1 {
			result.push('+');
			for _ in 0..shape[0] {
				result.push_str("-------------+");
			}
			result.push('\n');
			
			result.push('|');
			for x in 0..shape[0] {
				result.push_str(&format!(" Elected: {:2}", self[&[x]].elected));
				result.push_str(if x == 0 { " ‖" } else { " |" });
			}
			result.push('\n');
			
			result.push('|');
			for x in 0..shape[0] {
				result.push_str(&format!(" Min:     {:2}", self[&[x]].min));
				result.push_str(if x == 0 { " ‖" } else { " |" });
			}
			result.push('\n');
			
			result.push('|');
			for x in 0..shape[0] {
				result.push_str(&format!(" Max:     {:2}", self[&[x]].max));
				result.push_str(if x == 0 { " ‖" } else { " |" });
			}
			result.push('\n');
			
			result.push('|');
			for x in 0..shape[0] {
				result.push_str(&format!(" Cands:   {:2}", self[&[x]].cands));
				result.push_str(if x == 0 { " ‖" } else { " |" });
			}
			result.push('\n');
			
			result.push('+');
			for _ in 0..shape[0] {
				result.push_str("-------------+");
			}
			result.push('\n');
		} else if shape.len() == 2 {
			for y in 0..shape[1] {
				result.push('+');
				for _ in 0..shape[0] {
					result.push_str(if y == 1 { "=============+" } else { "-------------+" });
				}
				result.push('\n');
				
				result.push('|');
				for x in 0..shape[0] {
					result.push_str(&format!(" Elected: {:2}", self[&[x, y]].elected));
					result.push_str(if x == 0 { " ‖" } else { " |" });
				}
				result.push('\n');
				
				result.push('|');
				for x in 0..shape[0] {
					result.push_str(&format!(" Min:     {:2}", self[&[x, y]].min));
					result.push_str(if x == 0 { " ‖" } else { " |" });
				}
				result.push('\n');
				
				result.push('|');
				for x in 0..shape[0] {
					result.push_str(&format!(" Max:     {:2}", self[&[x, y]].max));
					result.push_str(if x == 0 { " ‖" } else { " |" });
				}
				result.push('\n');
				
				result.push('|');
				for x in 0..shape[0] {
					result.push_str(&format!(" Cands:   {:2}", self[&[x, y]].cands));
					result.push_str(if x == 0 { " ‖" } else { " |" });
				}
				result.push('\n');
			}
			
			result.push('+');
			for _ in 0..shape[0] {
				result.push_str("-------------+");
			}
			result.push('\n');
		} else {
			todo!();
		}
		
		return f.write_str(&result);
	}
}

impl ops::Index<&[usize]> for ConstraintMatrix {
	type Output = ConstraintMatrixCell;
	fn index(&self, index: &[usize]) -> &Self::Output { &self.0[index] }
}
impl ops::IndexMut<&[usize]> for ConstraintMatrix {
	fn index_mut(&mut self, index: &[usize]) -> &mut Self::Output { &mut self.0[index] }
}

/// Return the [Candidate]s referred to in the given [ConstraintMatrixCell] at location `idx`
fn candidates_in_constraint_cell<'a, N: Number>(election: &'a Election<N>, candidates: &CandidateMap<CountCard<N>>, idx: &[usize]) -> Vec<&'a Candidate> {
	let mut result: Vec<&Candidate> = Vec::new();
	for (i, candidate) in election.candidates.iter().enumerate() {
		let cc = &candidates[candidate];
		if cc.state != CandidateState::Hopeful {
			continue;
		}
		
		// Is this candidate within this constraint cell?
		let mut matches = true;
		for (coord, constraint) in idx.iter().zip(election.constraints.as_ref().unwrap().0.iter()) {
			let group = &constraint.groups[coord - 1]; // The group referred to by this constraint cell
			if !group.candidates.contains(&i) {
				matches = false;
				break;
			}
		}
		
		if matches {
			result.push(candidate);
		}
	}
	
	return result;
}

/// Clone and update the constraints matrix, with the state of the given candidates set to candidate_state – check if a conformant result is possible
pub fn test_constraints_any_time<N: Number>(state: &CountState<N>, candidates: &[&Candidate], candidate_state: CandidateState) -> Result<(), ConstraintError> {
	if state.constraint_matrix.is_none() {
		return Ok(());
	}
	let mut cm = state.constraint_matrix.as_ref().unwrap().clone();
	
	let mut trial_candidates = state.candidates.clone(); // TODO: Can probably be optimised by not cloning CountCard::parcels
	for candidate in candidates {
		trial_candidates.get_mut(candidate).unwrap().state = candidate_state;
	}
	
	// Update cands/elected
	cm.update_from_state(state.election, &trial_candidates);
	cm.recount_cands();
	
	// Iterate for stable state
	while !cm.step()? {}
	
	return Ok(());
}

/// Update the constraints matrix, and perform the necessary actions given by [STVOptions::constraint_mode]
pub fn update_constraints<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> bool {
	if state.constraint_matrix.is_none() {
		return false;
	}
	let cm = state.constraint_matrix.as_mut().unwrap();
	
	// Update cands/elected
	cm.update_from_state(state.election, &state.candidates);
	cm.recount_cands();
	
	// Iterate for stable state
	while !cm.step().expect("No conformant result is possible") {}
	
	if state.num_elected == state.election.seats {
		// Election is complete, so skip guarding/dooming candidates
		return false;
	}
	
	match opts.constraint_mode {
		ConstraintMode::GuardDoom => {
			// Check for guarded or doomed candidates
			let mut guarded_or_doomed = false;
			
			for idx in ndarray::indices(cm.0.shape()) {
				if (0..idx.ndim()).fold(0, |acc, d| if idx[d] == 0 { acc + 1 } else { acc }) != 0 {
					continue;
				}
				let cell = &cm.0[&idx];
				
				if cell.elected == cell.max {
					// Doom remaining candidates in this cell
					let doomed = candidates_in_constraint_cell(state.election, &state.candidates, idx.slice());
					if !doomed.is_empty() {
						for candidate in doomed.iter() {
							state.candidates.get_mut(candidate).unwrap().state = CandidateState::Doomed;
						}
						
						state.logger.log_smart(
							"{} must be doomed to comply with constraints.",
							"{} must be doomed to comply with constraints.",
							doomed.iter().map(|c| c.name.as_str()).sorted().collect()
						);
						
						guarded_or_doomed = true;
					}
				}
				if cell.cands == cell.min {
					// Guard remaining candidates in this cell
					let guarded = candidates_in_constraint_cell(state.election, &state.candidates, idx.slice());
					if !guarded.is_empty() {
						for candidate in guarded.iter() {
							state.candidates.get_mut(candidate).unwrap().state = CandidateState::Guarded;
						}
						
						state.logger.log_smart(
							"{} must be guarded to comply with constraints.",
							"{} must be guarded to comply with constraints.",
							guarded.iter().map(|c| c.name.as_str()).sorted().collect()
						);
						
						guarded_or_doomed = true;
					}
				}
			}
			
			return guarded_or_doomed;
		}
		ConstraintMode::RepeatCount => { return false; } // No action needed here: elect_hopefuls checks test_constraints_immediate
	}
}

// ----------------------------------
// FOR --constraint-mode repeat_count
// ----------------------------------

/// Check constraints, with the state of the given candidates set to candidate_state – check if this immediately violates constraints
pub fn test_constraints_immediate<'a, N: Number>(state: &CountState<'a, N>, candidates: &[&Candidate], candidate_state: CandidateState) -> Result<(), (&'a Constraint, &'a ConstrainedGroup)> {
	if state.election.constraints.is_none() {
		return Ok(());
	}
	
	let mut trial_candidates = state.candidates.clone(); // TODO: Can probably be optimised by not cloning CountCard::parcels
	for candidate in candidates {
		trial_candidates.get_mut(candidate).unwrap().state = candidate_state;
	}
	
	if let Some((a, b)) = state.election.constraints.as_ref().unwrap().exceeds_maximum(state.election, trial_candidates) {
		return Err((a, b));
	}
	
	return Ok(());
}

/// Initialise the [Election] as required for --constraint-mode repeat_count
pub fn init_repeat_count<N: Number>(election: &mut Election<N>) {
	// Add dummy candidates
	let mut new_candidates = Vec::new();
	for candidate in &election.candidates {
		let mut new_candidate = candidate.clone();
		new_candidate.index += election.candidates.len(); // Ensure unique index
		new_candidate.is_dummy = true;
		new_candidates.push(new_candidate);
	}
	election.candidates.append(&mut new_candidates);
}

/// Initialise the rollback for [ConstraintMode::RepeatCount]
pub fn init_repeat_count_rollback<'a, N: Number>(state: &mut CountState<'a, N>, constraint: &'a Constraint, group: &'a ConstrainedGroup) {
	let mut rollback_candidates = CandidateMap::with_capacity(state.candidates.len());
	let rollback_exhausted = state.exhausted.clone();
	
	// Copy ballot papers to rollback state
	for (candidate, count_card) in state.candidates.iter_mut() {
		rollback_candidates.insert(candidate, count_card.clone());
	}
	
	state.rollback_state = RollbackState::NeedsRollback { candidates: Some(rollback_candidates), exhausted: Some(rollback_exhausted), constraint, group };
}

/// Process one stage of rollback for [ConstraintMode::RepeatCount]
pub fn rollback_one_stage<N: Number>(state: &mut CountState<N>, opts: &STVOptions) -> Result<(), STVError>
where
	for<'r> &'r N: ops::Add<&'r N, Output=N>,
	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>
{
	if let RollbackState::NeedsRollback { candidates, exhausted, constraint, group } = &mut state.rollback_state {
		let mut candidates = candidates.take().unwrap();
		
		// Exclude candidates who cannot be elected due to constraint violations
		let order_excluded = state.num_excluded + 1;
		let mut excluded_candidates = Vec::new();
		for candidate_idx in &group.candidates {
			let count_card = state.candidates.get_mut(&state.election.candidates[*candidate_idx]).unwrap();
			if count_card.state == CandidateState::Hopeful {
				count_card.state = CandidateState::Excluded;
				count_card.finalised = true;
				state.num_excluded += 1;
				count_card.order_elected = -(order_excluded as isize);
				excluded_candidates.push(state.election.candidates[*candidate_idx].name.as_str());
			}
		}
		
		// Prepare dummy candidates, etc.
		for candidate in &state.election.candidates {
			if candidate.is_dummy {
				continue;
			}
			
			// Move ballot papers to dummy candidate
			let dummy_candidate = state.election.candidates.iter().find(|c| c.name == candidate.name && c.is_dummy).unwrap();
			let dummy_count_card = state.candidates.get_mut(dummy_candidate).unwrap();
			dummy_count_card.parcels.append(&mut candidates.get_mut(candidate).unwrap().parcels);
			dummy_count_card.votes = candidates[candidate].votes.clone();
			
			// Reset count
			let count_card = state.candidates.get_mut(candidate).unwrap();
			count_card.parcels.clear();
			count_card.votes = N::new();
			count_card.transfers = N::new();
			
			if candidates[candidate].state == CandidateState::Elected {
				if &candidates[candidate].votes > state.quota.as_ref().unwrap() {
					count_card.votes = state.quota.as_ref().unwrap().clone();
				} else {
					count_card.votes = candidates[candidate].votes.clone();
				}
			}
		}
		
		state.title = StageKind::Rollback;
		state.logger.log_smart(
			"Rolled back to apply constraints. {} is excluded.",
			"Rolled back to apply constraints. {} are excluded.",
			excluded_candidates.into_iter().sorted().collect()
		);
		
		state.rollback_state = RollbackState::RollingBack { candidates: Some(candidates), exhausted: exhausted.take(), candidate_distributing: None, constraint: Some(constraint), group: Some(group) };
		return Ok(());
	}
	
	if let RollbackState::RollingBack { candidates, exhausted, candidate_distributing, constraint, group } = &mut state.rollback_state {
		let candidates = candidates.take().unwrap();
		let mut exhausted = exhausted.take().unwrap();
		let mut candidate_distributing = candidate_distributing.take();
		let constraint = constraint.take().unwrap();
		let group = group.take().unwrap();
		
		// --------------------
		// Distribute surpluses
		
		let has_surplus: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of tie
			.filter(|c| {
				let cc = &candidates[c];
				if !c.is_dummy && cc.state == CandidateState::Elected && !cc.finalised {
					let dummy_candidate = state.election.candidates.iter().find(|x| x.name == c.name && x.is_dummy).unwrap();
					!state.candidates[dummy_candidate].finalised
				} else {
					false
				}
			})
			.collect();
		
		if !has_surplus.is_empty() {
			// Distribute top candidate's surplus
			let max_cands = match opts.surplus_order {
				SurplusOrder::BySize => {
					ties::multiple_max_by(&has_surplus, |c| &candidates[c].votes)
				}
				SurplusOrder::ByOrder => {
					ties::multiple_min_by(&has_surplus, |c| candidates[c].order_elected)
				}
			};
			let elected_candidate = if max_cands.len() > 1 {
				stv::choose_highest(state, opts, &max_cands, "Which candidate's surplus to distribute?")?
			} else {
				max_cands[0]
			};
			
			let dummy_candidate = state.election.candidates.iter().find(|c| c.name == elected_candidate.name && c.is_dummy).unwrap();
			
			match opts.surplus {
				SurplusMethod::WIG | SurplusMethod::UIG | SurplusMethod::EG => { gregory::distribute_surplus(state, opts, dummy_candidate); }
				SurplusMethod::IHare | SurplusMethod::Hare => { sample::distribute_surplus(state, opts, dummy_candidate)?; }
				_ => unreachable!()
			}
			
			state.rollback_state = RollbackState::RollingBack { candidates: Some(candidates), exhausted: Some(exhausted), candidate_distributing, constraint: Some(constraint), group: Some(group) };
			rollback_check_complete(state);
			return Ok(());
		}
		
		// ----------------------------------
		// Distribute exhausted ballot papers
		// FIXME: Untested!
		
		if exhausted.parcels.iter().any(|p| !p.votes.is_empty()) {
			// Use arbitrary dummy candidate
			let dummy_candidate = state.election.candidates.iter().find(|c| c.is_dummy).unwrap();
			
			let dummy_count_card = state.candidates.get_mut(dummy_candidate).unwrap();
			dummy_count_card.parcels.append(&mut exhausted.parcels);
			
			// Nasty hack to check if continuing distribution!!
			if let StageKind::RollbackExhausted = state.title {
				// Continuing
				state.logger.log_literal(String::from("Continuing distribution of exhausted ballots."));
			} else {
				state.title = StageKind::RollbackExhausted;
				state.logger.log_literal(String::from("Distributing exhausted ballots."));
			}
			
			stv::exclude_candidates(state, opts, vec![dummy_candidate], "Distribution")?;
			
			let dummy_count_card = state.candidates.get_mut(dummy_candidate).unwrap();
			exhausted.parcels.append(&mut dummy_count_card.parcels);
			
			state.rollback_state = RollbackState::RollingBack { candidates: Some(candidates), exhausted: Some(exhausted), candidate_distributing: None, constraint: Some(constraint), group: Some(group) };
			rollback_check_complete(state);
			return Ok(());
		}
		
		// ------------------------------------------
		// Distribute ballots of electable candidates
		
		let electable_candidates_old: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of multiple
			.filter(|c| {
				let cc = &candidates[c];
				let cand_idx = state.election.candidates.iter().position(|x| x == *c).unwrap();
				if !c.is_dummy && !group.candidates.contains(&cand_idx) && !cc.finalised {
					let dummy_candidate = state.election.candidates.iter().find(|x| x.name == c.name && x.is_dummy).unwrap();
					!state.candidates[dummy_candidate].finalised
				} else {
					false
				}
			})
			.collect();
		
		if !electable_candidates_old.is_empty() {
			if candidate_distributing.is_none() || !electable_candidates_old.contains(candidate_distributing.as_ref().unwrap()) {
				if electable_candidates_old.len() > 1 {
					// Determine or prompt for which candidate to distribute
					for strategy in opts.ties.iter() {
						match strategy {
							TieStrategy::Random(_) | TieStrategy::Prompt => {
								candidate_distributing = Some(strategy.choose_lowest(state, opts, &electable_candidates_old, "Which candidate's ballots to distribute?").unwrap());
								break;
							}
							TieStrategy::Forwards | TieStrategy::Backwards => {}
						}
					}
				} else {
					candidate_distributing = Some(electable_candidates_old[0]);
				}
				
				state.logger.log_smart(
					"Distributing ballots of {}.",
					"Distributing ballots of {}.",
					vec![candidate_distributing.unwrap().name.as_str()]
				);
			} else {
				state.logger.log_smart(
					"Continuing distribution of ballots of {}.",
					"Continuing distribution of ballots of {}.",
					vec![candidate_distributing.unwrap().name.as_str()]
				);
			}
			
			let candidate_distributing = candidate_distributing.unwrap();
			let dummy_candidate = state.election.candidates.iter().find(|c| c.name == candidate_distributing.name && c.is_dummy).unwrap();
			
			state.title = StageKind::BallotsOf(candidate_distributing);
			stv::exclude_candidates(state, opts, vec![dummy_candidate], "Distribution")?;
			
			state.rollback_state = RollbackState::RollingBack { candidates: Some(candidates), exhausted: Some(exhausted), candidate_distributing: Some(candidate_distributing), constraint: Some(constraint), group: Some(group) };
			rollback_check_complete(state);
			return Ok(());
		}
		
		// --------------------------------------------
		// Distribute ballots of unelectable candidates
		
		let unelectable_candidates_old: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of multiple
			.filter(|c| {
				let cc = &candidates[c];
				let cand_idx = state.election.candidates.iter().position(|x| x == *c).unwrap();
				if !c.is_dummy && group.candidates.contains(&cand_idx) && !cc.finalised {
					let dummy_candidate = state.election.candidates.iter().find(|x| x.name == c.name && x.is_dummy).unwrap();
					!state.candidates[dummy_candidate].finalised
				} else {
					false
				}
			})
			.collect();
		
		if !unelectable_candidates_old.is_empty() {
			if candidate_distributing.is_none() || !unelectable_candidates_old.contains(candidate_distributing.as_ref().unwrap()) {
				if unelectable_candidates_old.len() > 1 {
					// Determine or prompt for which candidate to distribute
					for strategy in opts.ties.iter() {
						match strategy {
							TieStrategy::Random(_) | TieStrategy::Prompt => {
								candidate_distributing = Some(strategy.choose_lowest(state, opts, &unelectable_candidates_old, "Which candidate's ballots to distribute?").unwrap());
								break;
							}
							TieStrategy::Forwards | TieStrategy::Backwards => {}
						}
					}
				} else {
					candidate_distributing = Some(unelectable_candidates_old[0]);
				}
				
				state.logger.log_smart(
					"Distributing ballots of {}.",
					"Distributing ballots of {}.",
					vec![candidate_distributing.unwrap().name.as_str()]
				);
			} else {
				state.logger.log_smart(
					"Continuing distribution of ballots of {}.",
					"Continuing distribution of ballots of {}.",
					vec![candidate_distributing.unwrap().name.as_str()]
				);
			}
			
			let candidate_distributing = candidate_distributing.unwrap();
			let dummy_candidate = state.election.candidates.iter().find(|c| c.name == candidate_distributing.name && c.is_dummy).unwrap();
			
			state.title = StageKind::BallotsOf(candidate_distributing);
			stv::exclude_candidates(state, opts, vec![dummy_candidate], "Distribution")?;
			
			state.rollback_state = RollbackState::RollingBack { candidates: Some(candidates), exhausted: Some(exhausted), candidate_distributing: Some(candidate_distributing), constraint: Some(constraint), group: Some(group) };
			rollback_check_complete(state);
			return Ok(());
		}
	}
	
	unreachable!();
}

fn rollback_check_complete<N: Number>(state: &mut CountState<N>) {
	if let RollbackState::RollingBack { candidates, exhausted, candidate_distributing: _, constraint: _, group: _ } = &state.rollback_state {
		let candidates = candidates.as_ref().unwrap();
		let exhausted = exhausted.as_ref().unwrap();
		
		let has_surplus: Vec<&Candidate> = state.election.candidates.iter() // Present in order in case of tie
			.filter(|c| {
				let cc = &candidates[c];
				if !c.is_dummy && cc.state == CandidateState::Elected && !cc.finalised {
					let dummy_candidate = state.election.candidates.iter().find(|x| x.name == c.name && x.is_dummy).unwrap();
					!state.candidates[dummy_candidate].finalised
				} else {
					false
				}
			})
			.collect();
		
		if !has_surplus.is_empty() {
			return;
		}
		
		if exhausted.parcels.iter().any(|p| !p.votes.is_empty()) {
			return;
		}
		
		let continuing_candidates: Vec<&Candidate> = state.election.candidates.iter()
			.filter(|c| {
				let cc = &candidates[c];
				if !c.is_dummy && !cc.finalised {
					let dummy_candidate = state.election.candidates.iter().find(|x| x.name == c.name && x.is_dummy).unwrap();
					!state.candidates[dummy_candidate].finalised
				} else {
					false
				}
			})
			.collect();
		
		if !continuing_candidates.is_empty() {
			return;
		}
		
		// ---------------------------
		// Rollback complete: finalise
		
		// Delete dummy candidates
		for (candidate, count_card) in state.candidates.iter_mut() {
			if candidate.is_dummy {
				count_card.state = CandidateState::Withdrawn;
				count_card.parcels.clear();
				count_card.votes = N::new();
				count_card.transfers = N::new();
			}
		}
		
		state.logger.log_literal(String::from("Rollback complete."));
		state.rollback_state = RollbackState::Normal;
		state.num_excluded = state.candidates.values().filter(|cc| cc.state == CandidateState::Excluded).count();
		
		return;
	}
	
	unreachable!();
}

#[cfg(test)]
mod tests {
	use super::*;
	
	fn assert_cell(cm: &ConstraintMatrix, idx: &[usize], elected: usize, min: usize, max: usize, cands: usize) {
		assert_eq!(cm[idx].elected, elected, "Failed to validate elected at {:?}", idx);
		assert_eq!(cm[idx].min, min, "Failed to validate min at {:?}", idx);
		assert_eq!(cm[idx].max, max, "Failed to validate min at {:?}", idx);
		assert_eq!(cm[idx].cands, cands, "Failed to validate cands at {:?}", idx);
	}
	
	#[test]
	fn cm_otten() {
		let mut cm = ConstraintMatrix::new(&mut [3, 2]);
		
		// Totals
		let c = &mut cm[&[0, 1]]; c.min = 7; c.max = 7;
		let c = &mut cm[&[0, 2]]; c.min = 7; c.max = 7;
		let c = &mut cm[&[1, 0]]; c.min = 7; c.max = 7;
		let c = &mut cm[&[2, 0]]; c.min = 6; c.max = 6;
		let c = &mut cm[&[3, 0]]; c.min = 1; c.max = 1;
		
		// Candidates
		let c = &mut cm[&[1, 1]]; c.cands = 4;
		let c = &mut cm[&[2, 1]]; c.cands = 11;
		let c = &mut cm[&[3, 1]]; c.cands = 2;
		let c = &mut cm[&[1, 2]]; c.cands = 7;
		let c = &mut cm[&[2, 2]]; c.cands = 3;
		let c = &mut cm[&[3, 2]]; c.cands = 1;
		
		// Init
		cm.init();
		while !cm.step().expect("No conformant result") {}
		println!("{}", cm);
		
		assert_cell(&cm, &[1, 1], 0, 0, 4, 4);
		assert_cell(&cm, &[2, 1], 0, 3, 6, 11);
		assert_cell(&cm, &[3, 1], 0, 0, 1, 2);
		assert_cell(&cm, &[0, 1], 0, 7, 7, 17);
		assert_cell(&cm, &[1, 2], 0, 3, 7, 7);
		assert_cell(&cm, &[2, 2], 0, 0, 3, 3);
		assert_cell(&cm, &[3, 2], 0, 0, 1, 1);
		assert_cell(&cm, &[0, 2], 0, 7, 7, 11);
		assert_cell(&cm, &[1, 0], 0, 7, 7, 11);
		assert_cell(&cm, &[2, 0], 0, 6, 6, 14);
		assert_cell(&cm, &[3, 0], 0, 1, 1, 3);
		assert_cell(&cm, &[0, 0], 0, 14, 14, 28);
		
		// Election of Welsh man
		cm[&[3, 1]].elected += 1;
		cm.recount_cands();
		while !cm.step().expect("No conformant result") {}
		println!("{}", cm);
		
		assert_cell(&cm, &[1, 1], 0, 0, 3, 4);
		assert_cell(&cm, &[2, 1], 0, 3, 6, 11);
		assert_cell(&cm, &[3, 1], 1, 1, 1, 2);
		assert_cell(&cm, &[0, 1], 1, 7, 7, 17); // Error in Otten paper
		assert_cell(&cm, &[1, 2], 0, 4, 7, 7);
		assert_cell(&cm, &[2, 2], 0, 0, 3, 3);
		assert_cell(&cm, &[3, 2], 0, 0, 0, 1);
		assert_cell(&cm, &[0, 2], 0, 7, 7, 11);
		assert_cell(&cm, &[1, 0], 0, 7, 7, 11);
		assert_cell(&cm, &[2, 0], 0, 6, 6, 14);
		assert_cell(&cm, &[3, 0], 1, 1, 1, 3);
		assert_cell(&cm, &[0, 0], 1, 14, 14, 28);
		
		// Remaining Welsh man, Welsh woman doomed
		cm[&[3, 1]].cands -= 1;
		cm[&[3, 2]].cands -= 1;
		
		// Election of 2 English men, 2 English women
		// Exclusion of 1 Scottish woman
		cm[&[1, 1]].elected += 2;
		cm[&[1, 2]].elected += 2;
		cm[&[2, 2]].cands -= 1;
		cm.recount_cands();
		while !cm.step().expect("No conformant result") {}
		println!("{}", cm);
		
		assert_cell(&cm, &[1, 1], 2, 2, 2, 4);
		assert_cell(&cm, &[2, 1], 0, 4, 4, 11);
		assert_cell(&cm, &[3, 1], 1, 1, 1, 1);
		assert_cell(&cm, &[0, 1], 3, 7, 7, 16); // Error in Otten paper
		assert_cell(&cm, &[1, 2], 2, 5, 5, 7);
		assert_cell(&cm, &[2, 2], 0, 2, 2, 2);
		assert_cell(&cm, &[3, 2], 0, 0, 0, 0);
		assert_cell(&cm, &[0, 2], 2, 7, 7, 9);
		assert_cell(&cm, &[1, 0], 4, 7, 7, 11);
		assert_cell(&cm, &[2, 0], 0, 6, 6, 13);
		assert_cell(&cm, &[3, 0], 1, 1, 1, 1);
		assert_cell(&cm, &[0, 0], 5, 14, 14, 25);
	}
}