/*  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 super::STVOptions;

use crate::election::{Candidate, CountState};
use crate::numbers::Number;

use std::collections::HashMap;

/// Table describing vote transfers during a surplus distribution or exclusion
pub struct TransferTable<'e, N: Number> {
	/// Columns in the table
	pub columns: Vec<TransferTableColumn<'e, N>>,
}

impl<'e, N: Number> TransferTable<'e, N> {
	/// Return a new [TransferTable]
	pub fn new() -> Self {
		TransferTable {
			columns: Vec::new(),
		}
	}
	
	/// Record the specified transfer
	pub fn add_transfers(&mut self, value_fraction: &N, candidate: &'e Candidate, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction {
				col.add_transfers(candidate, ballots);
				return;
			}
		}
		
		let mut col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			cells: HashMap::new(),
			exhausted: TransferTableCell { ballots: N::new() },
		};
		col.add_transfers(candidate, ballots);
		self.columns.push(col);
	}
	
	/// Record the specified exhaustion
	pub fn add_exhausted(&mut self, value_fraction: &N, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction {
				col.exhausted.ballots += ballots;
				return;
			}
		}
		
		let col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			cells: HashMap::new(),
			exhausted: TransferTableCell { ballots: ballots.clone() },
		};
		self.columns.push(col);
	}
	
	/// Apply the transfers described in the table to the count sheet
	///
	/// Credit continuing candidates and exhausted pile with the appropriate number of ballot papers and votes.
	pub fn apply_to(&self, state: &mut CountState<N>, opts: &STVOptions) -> N {
		// TODO: SumSurplusTransfers
		
		let mut checksum = N::new();
		
		// Credit transferred votes
		for (candidate, count_card) in state.candidates.iter_mut() {
			let mut votes_transferred = N::new();
			let mut ballots_transferred = N::new();
			
			for column in self.columns.iter() {
				if let Some(cell) = column.cells.get(*candidate) {
					votes_transferred += cell.ballots.clone() * &column.value_fraction;
					ballots_transferred += &cell.ballots;
				}
			}
			
			if let Some(dps) = opts.round_votes {
				votes_transferred.floor_mut(dps);
			}
			
			count_card.transfer(&votes_transferred);
			count_card.ballot_transfers += ballots_transferred;
			
			checksum += votes_transferred;
		}
		
		// Credit exhausted votes
		let mut votes_transferred = N::new();
		let mut ballots_transferred = N::new();
		
		for column in self.columns.iter() {
			votes_transferred += column.exhausted.ballots.clone() * &column.value_fraction;
			ballots_transferred += &column.exhausted.ballots;
		}
		
		if let Some(dps) = opts.round_votes {
			votes_transferred.floor_mut(dps);
		}
		
		state.exhausted.transfer(&votes_transferred);
		state.exhausted.ballot_transfers += ballots_transferred;
		
		checksum += votes_transferred;
		
		return checksum;
	}
}

/// Column in a [TransferTable]
pub struct TransferTableColumn<'e, N: Number> {
	/// Value fraction of ballots counted in this column
	pub value_fraction: N,
	
	/// Cells in this column
	pub cells: HashMap<&'e Candidate, TransferTableCell<N>>,
	
	/// Exhausted cell
	pub exhausted: TransferTableCell<N>,
}

impl<'e, N: Number> TransferTableColumn<'e, N> {
	/// Record the specified transfer
	pub fn add_transfers(&mut self, candidate: &'e Candidate, ballots: &N) {
		if let Some(cell) = self.cells.get_mut(candidate) {
			cell.ballots += ballots;
		} else {
			let cell = TransferTableCell {
				ballots: ballots.clone(),
			};
			self.cells.insert(candidate, cell);
		}
	}
}

/// Cell in a [TransferTable], representing transfers to one candidate at a particular value
pub struct TransferTableCell<N: Number> {
	/// Ballots transferred to this candidate
	pub ballots: N,
}