/*  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 ndarray::{Array, Dimension, IxDyn};

use std::fmt;
use std::ops;

#[derive(Debug)]
enum ConstraintError {
	NoConformantResult,
}

#[derive(Clone)]
struct ConstraintMatrixCell {
	elected: usize,
	min: usize,
	max: usize,
	cands: usize,
}

struct ConstraintMatrix(Array<ConstraintMatrixCell, IxDyn>);

impl ConstraintMatrix {
	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,
			}
		));
	}
	
	pub fn init(&mut self) {
		let indices: Vec<IxDyn> = ndarray::indices(self.0.shape()).into_iter().collect();
		
		// 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 indices.iter() {
			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
	}
	
	pub fn recount_cands(&mut self) {
		let shape = Vec::from(self.0.shape());
		let indices: Vec<IxDyn> = ndarray::indices(self.0.shape()).into_iter().collect();
		
		// Compute cands/elected totals
		for nzeroes in 1..self.0.ndim()+1 {
			for idx in indices.iter() {
				// 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()).filter(|d| idx[*d] == 0).next().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;
				}
			}
		}
	}
	
	pub fn step(&mut self) -> Result<bool, ConstraintError> {
		let shape = Vec::from(self.0.shape());
		let indices: Vec<IxDyn> = ndarray::indices(self.0.shape()).into_iter().collect();
		
		for idx in indices.iter() {
			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 inner 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
		for y in 0..shape[1] {
			result.push_str("+");
			for _ in 0..shape[0] {
				result.push_str(if y == 1 { "=============+" } else { "-------------+" });
			}
			result.push_str("\n");
			
			result.push_str("|");
			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_str("\n");
			
			result.push_str("|");
			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_str("\n");
			
			result.push_str("|");
			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_str("\n");
			
			result.push_str("|");
			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_str("\n");
		}
		
		result.push_str("+");
		for _ in 0..shape[0] {
			result.push_str("-------------+");
		}
		result.push_str("\n");
		
		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] }
}

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