/*  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, SumSurplusTransfersMode};

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, surplus: Option<&N>, surplus_numer: &Option<N>, surplus_denom: &Option<N>) -> N {
		// Use weighted rules if exclusion or WIGM
		let is_weighted = surplus.is_none() || opts.surplus.is_weighted();
		
		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();
			
			// If exclusion, or surplus at present value, or SumSurplusTransfersMode::ByValue
			if surplus_numer.is_none() || opts.sum_surplus_transfers == SumSurplusTransfersMode::ByValue {
				// Calculate transfer across all votes in this parcel
				for column in self.columns.iter() {
					if let Some(cell) = column.cells.get(*candidate) {
						if is_weighted {
							votes_transferred += cell.ballots.clone() * &column.value_fraction;
						}
						ballots_transferred += &cell.ballots;
					}
				}
				
				if !is_weighted {
					votes_transferred = ballots_transferred.clone();
				}
				
				// If surplus, multiply by surplus fraction
				if let Some(n) = &surplus_numer {
					votes_transferred *= n;
				}
				if let Some(n) = &surplus_denom {
					votes_transferred /= n;
				}
			} else if opts.sum_surplus_transfers == SumSurplusTransfersMode::PerBallot {
				// Sum transfer per each individual ballot
				for column in self.columns.iter() {
					if let Some(cell) = column.cells.get(*candidate) {
						ballots_transferred += &cell.ballots;
						
						let mut new_value_fraction;
						if is_weighted {
							new_value_fraction = column.value_fraction.clone();
							// If surplus, multiply by surplus fraction
							if let Some(n) = &surplus_numer {
								new_value_fraction *= n;
							}
							if let Some(n) = &surplus_denom {
								new_value_fraction /= n;
							}
							// Round if required
							if let Some(dps) = opts.round_values {
								new_value_fraction.floor_mut(dps);
							}
						} else {
							if let Some(n) = &surplus_numer {
								new_value_fraction = n.clone();
							} else {
								// Transferred at original value
								new_value_fraction = column.value_fraction.clone();
							}
							if let Some(n) = &surplus_denom {
								new_value_fraction /= n;
							}
							// Round if required
							if let Some(dps) = opts.round_values {
								new_value_fraction.floor_mut(dps);
							}
						}
						
						votes_transferred += cell.ballots.clone() * new_value_fraction;
					}
				}
			} else {
				unreachable!();
			}
			
			// Round if required
			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
		// If exclusion or not --transferable-only
		if surplus.is_none() || !opts.transferable_only {
			// Standard rules
			let mut votes_transferred = N::new();
			let mut ballots_transferred = N::new();
			
			for column in self.columns.iter() {
				if is_weighted {
					votes_transferred += column.exhausted.ballots.clone() * &column.value_fraction;
				}
				ballots_transferred += &column.exhausted.ballots;
			}
			
			if !is_weighted {
				votes_transferred = ballots_transferred.clone();
			}
			
			// If surplus, multiply by surplus fraction
			if let Some(n) = &surplus_numer {
				votes_transferred *= n;
			}
			if let Some(n) = &surplus_denom {
				votes_transferred /= n;
			}
			
			// Round if required
			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;
		} else {
			// Credit only nontransferable difference
			if surplus_numer.is_none() {
				// TODO: Is there a purer way of calculating this?
				let difference = surplus.unwrap().clone() - &checksum;
				state.exhausted.transfer(&difference);
				checksum += difference;
				
				for column in self.columns.iter() {
					state.exhausted.ballot_transfers += &column.exhausted.ballots;
				}
			} else {
				// No ballots exhaust
			}
		}
		
		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,
}