/*  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::election::{Candidate, CandidateState, CountCard, CountState, Election};
use crate::numbers::Number;
use crate::stv::{ConstraintMode, STVOptions};

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

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

use std::collections::HashMap;
use std::fmt;
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<I: Iterator<Item=String>>(lines: I) -> Self {
		let mut constraints = Constraints(Vec::new());
		
		for line in lines {
			let mut bits = line.split(' ').peekable();
			
			// Read constraint category and group
			let constraint_name = read_quoted_string(&mut bits);
			let group_name = read_quoted_string(&mut bits);
			
			// Read min, max
			let min: usize = bits.next().expect("Syntax Error").parse().expect("Syntax Error");
			let max: usize = bits.next().expect("Syntax Error").parse().expect("Syntax Error");
			
			// Read candidates
			let mut candidates: Vec<usize> = Vec::new();
			for x in bits {
				candidates.push(x.parse::<usize>().expect("Syntax Error") - 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) {
				panic!("Duplicate group \"{}\" in constraint \"{}\"", group_name, constraint.name);
			}
			
			constraint.groups.push(ConstrainedGroup {
				name: group_name,
				candidates,
				min,
				max,
			});
		}
		
		// TODO: Validate constraints
		
		return constraints;
	}
}

/// Read an optionally quoted string, returning the string without quotes
fn read_quoted_string<'a, I: Iterator<Item=&'a str>>(bits: &mut I) -> String {
	let x = bits.next().expect("Syntax Error");
	if let Some(x1) = x.strip_prefix('"') {
		if let Some(x2) = x.strip_suffix('"') {
			// Complete string
			return String::from(x2);
		} else {
			// Incomplete string
			let mut result = String::from(x1);
			
			// Read until matching "
			loop {
				let x = bits.next().expect("Syntax Error");
				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 result;
		}
	} else {
		// Unquoted string
		return 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,
}

/// 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: &HashMap<&Candidate, 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() {
			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;
					}
				}
				panic!("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: &HashMap<&Candidate, 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
pub fn try_constraints<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;
		}
		_ => { todo!() }
	}
	
	//return false;
}

#[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);
	}
}