/*  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, CountState};
use crate::numbers::Number;
use crate::stv::{STVOptions, RoundSubtransfersMode};

#[cfg(not(target_arch = "wasm32"))]
use prettytable::{Cell, Row, Table};
#[cfg(target_arch = "wasm32")]
use super::prettytable_html::{Cell, Row, Table};

use std::cmp::max;

/// Table describing vote transfers during a surplus distribution or exclusion
#[derive(Clone)]
pub struct TransferTable<'e, N: Number> {
	/// Continuing candidates
	pub hopefuls: Vec<&'e Candidate>,
	
	/// Columns in the table
	pub columns: Vec<TransferTableColumn<'e, N>>,
	
	/// Total column
	pub total: TransferTableColumn<'e, N>,
	
	/// Size of surplus, or `None` if an exclusion
	pub surplus: Option<N>,
	/// Surplus fraction, or `None` if votes not reweighted/an exclusion (for display/optimisation only)
	pub surpfrac: Option<N>,
	/// Numerator of surplus fraction, or `None` if votes not reweighted/an exclusion
	pub surpfrac_numer: Option<N>,
	/// Denominator of surplus fraction, or `None`
	pub surpfrac_denom: Option<N>,
}

impl<'e, N: Number> TransferTable<'e, N> {
	/// Return a new [TransferTable] for an exclusion
	pub fn new_exclusion(hopefuls: Vec<&'e Candidate>) -> Self {
		let num_hopefuls = hopefuls.len();
		
		return TransferTable {
			hopefuls,
			columns: Vec::new(),
			total: TransferTableColumn {
				value_fraction: N::new(),
				order: 0,
				cells: CandidateMap::with_capacity(num_hopefuls),
				exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
				total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			},
			surplus: None,
			surpfrac: None,
			surpfrac_numer: None,
			surpfrac_denom: None,
		}
	}
	
	/// Return a new [TransferTable] for a surplus distribution
	pub fn new_surplus(hopefuls: Vec<&'e Candidate>, surplus: N, surpfrac: Option<N>, surpfrac_numer: Option<N>, surpfrac_denom: Option<N>) -> Self {
		let num_hopefuls = hopefuls.len();
		
		return TransferTable {
			hopefuls,
			columns: Vec::new(),
			total: TransferTableColumn {
				value_fraction: N::new(),
				order: 0,
				cells: CandidateMap::with_capacity(num_hopefuls),
				exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
				total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			},
			surplus: Some(surplus),
			surpfrac,
			surpfrac_numer,
			surpfrac_denom,
		}
	}
	
	/// Record the specified transfer
	///
	/// order: Pass `None` to force a new column
	pub fn add_transfers(&mut self, value_fraction: &N, order: Option<usize>, candidate: &'e Candidate, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction && order.map(|o| col.order == o).unwrap_or(false) {
				col.add_transfers(candidate, ballots);
				return;
			}
		}
		
		let mut col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			order: order.unwrap_or(0),
			cells: CandidateMap::new(),
			exhausted: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
			total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
		};
		col.add_transfers(candidate, ballots);
		self.columns.push(col);
	}
	
	/// Record the specified exhaustion
	///
	/// order: Pass `None` to force a new column
	pub fn add_exhausted(&mut self, value_fraction: &N, order: Option<usize>, ballots: &N) {
		for col in self.columns.iter_mut() {
			if &col.value_fraction == value_fraction && order.map(|o| col.order == o).unwrap_or(false) {
				col.exhausted.ballots += ballots;
				return;
			}
		}
		
		let col = TransferTableColumn {
			value_fraction: value_fraction.clone(),
			order: order.unwrap_or(0),
			cells: CandidateMap::new(),
			exhausted: TransferTableCell { ballots: ballots.clone(), votes_in: N::new(), votes_out: N::new() },
			total: TransferTableCell { ballots: N::new(), votes_in: N::new(), votes_out: N::new() },
		};
		self.columns.push(col);
	}
	
	/// Calculate the votes to be transferred according to this table
	pub fn calculate(&mut self, opts: &STVOptions) {
		// Use weighted rules if exclusion or WIGM
		let is_weighted = self.surplus.is_none() || opts.surplus.is_weighted();
		
		// Iterate through columns
		// Sum votes_in, etc.
		for column in self.columns.iter_mut() {
			// Candidate votes
			for (candidate, cell) in column.cells.iter_mut() {
				column.total.ballots += &cell.ballots;
				self.total.add_transfers(candidate, &cell.ballots);
				self.total.total.ballots += &cell.ballots;
				
				let votes_in = cell.ballots.clone() * &column.value_fraction;
				cell.votes_in += &votes_in;
				column.total.votes_in += &votes_in;
				self.total.cells.get_mut(candidate).unwrap().votes_in += &votes_in;
				self.total.total.votes_in += votes_in;
			}
			
			// Exhausted votes
			column.total.ballots += &column.exhausted.ballots;
			self.total.exhausted.ballots += &column.exhausted.ballots;
			self.total.total.ballots += &column.exhausted.ballots;
			
			let votes_in = column.exhausted.ballots.clone() * &column.value_fraction;
			column.exhausted.votes_in += &votes_in;
			column.total.votes_in += &votes_in;
			self.total.exhausted.votes_in += &votes_in;
			self.total.total.votes_in += votes_in;
		}
		
		match opts.round_subtransfers {
			RoundSubtransfersMode::SingleStep => {
				// No need to calculate votes_out for each column
				
				// Calculate total votes_out per candidate
				for (_candidate, cell) in self.total.cells.iter_mut() {
					if is_weighted {
						// Weighted rules
						// Multiply votes in by surplus fraction
						cell.votes_out = multiply_surpfrac(cell.votes_in.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
					} else if self.surpfrac.is_none() {
						// Unweighted rules but transfer at values received
						cell.votes_out = cell.votes_in.clone();
					} else {
						// Unweighted rules
						// Multiply ballots in by surplus fraction
						cell.votes_out = multiply_surpfrac(cell.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
					}
					
					// Round if required
					if let Some(dps) = opts.round_votes {
						cell.votes_out.floor_mut(dps);
					}
				}
				
				// Calculate total exhausted votes
				if is_weighted {
					// Weighted rules
					// Multiply votes in by surplus fraction
					self.total.exhausted.votes_out = multiply_surpfrac(self.total.exhausted.votes_in.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
				} else if self.surpfrac.is_none() {
					// Unweighted rules but transfer at values received
					// This can only happen with --transferable-only, so this will be calculated in apply_to
				} else {
					// Unweighted rules
					// Multiply ballots in by surplus fraction
					self.total.exhausted.votes_out = multiply_surpfrac(self.total.exhausted.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
				}
				
				// Round if required
				if let Some(dps) = opts.round_votes {
					self.total.exhausted.votes_out.floor_mut(dps);
				}
			}
			RoundSubtransfersMode::ByValue | RoundSubtransfersMode::ByValueAndSource | RoundSubtransfersMode::ByParcel => {
				// Calculate votes_out for each column
				for column in self.columns.iter_mut() {
					// Calculate votes_out per candidate in the column
					for (_candidate, cell) in column.cells.iter_mut() {
						if is_weighted {
							// Weighted rules
							// Multiply votes in by surplus fraction
							cell.votes_out = multiply_surpfrac(cell.votes_in.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
						} else if self.surpfrac.is_none() {
							// Unweighted rules but transfer at values received
							cell.votes_out = cell.votes_in.clone();
						} else {
							// Unweighted rules
							// Multiply ballots in by surplus fraction
							cell.votes_out = multiply_surpfrac(cell.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
						}
						
						// Round if required
						if let Some(dps) = opts.round_votes {
							cell.votes_out.floor_mut(dps);
						}
					}
					
					// Calculate exhausted votes in the column
					if is_weighted {
						// Weighted rules
						// Multiply votes in by surplus fraction
						column.exhausted.votes_out = multiply_surpfrac(column.exhausted.votes_in.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
					} else if self.surpfrac.is_none() {
						// Unweighted rules but transfer at values received
						// This can only happen with --transferable-only, so this will be calculated in apply_to
					} else {
						// Unweighted rules
						// Multiply ballots in by surplus fraction
						column.exhausted.votes_out = multiply_surpfrac(column.exhausted.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
					}
					
					// Round if required
					if let Some(dps) = opts.round_votes {
						column.exhausted.votes_out.floor_mut(dps);
					}
				}
				
				// Sum total votes_out per candidate
				for (candidate, cell) in self.total.cells.iter_mut() {
					cell.votes_out = self.columns.iter().fold(N::new(), |mut acc, col| {
						if let Some(cell) = col.cells.get(candidate) {
							acc += &cell.votes_out
						}
						acc
					});
				}
				
				// Sum total exhausted votes
				self.total.exhausted.votes_out = self.columns.iter().fold(N::new(), |mut acc, col| { acc += &col.exhausted.votes_out; acc });
			}
			RoundSubtransfersMode::PerBallot => {
				// Calculate votes_out for each column
				for column in self.columns.iter_mut() {
					// Calculate votes_out per candidate in the column
					for (_candidate, cell) in column.cells.iter_mut() {
						if is_weighted {
							// Weighted rules
							// Multiply ballots in by new value fraction
							let mut new_value_fraction = multiply_surpfrac(column.value_fraction.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
							if let Some(dps) = opts.round_values {
								new_value_fraction.floor_mut(dps);
							}
							
							cell.votes_out = cell.ballots.clone() * new_value_fraction;
						} else if self.surpfrac.is_none() {
							// Unweighted rules but transfer at values received
							cell.votes_out = cell.votes_in.clone();
						} else {
							// Unweighted rules
							// Multiply ballots in by surplus fraction
							cell.votes_out = multiply_surpfrac(cell.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
						}
						
						// Round if required
						if let Some(dps) = opts.round_votes {
							cell.votes_out.floor_mut(dps);
						}
					}
					
					// Calculate exhausted votes in the column
					if is_weighted {
						// Weighted rules
						// Multiply ballots in by new value fraction
						let mut new_value_fraction = multiply_surpfrac(column.value_fraction.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
						if let Some(dps) = opts.round_values {
							new_value_fraction.floor_mut(dps);
						}
						
						column.exhausted.votes_out = column.exhausted.ballots.clone() * new_value_fraction;
					} else if self.surpfrac.is_none() {
						// Unweighted rules but transfer at values received
						// This can only happen with --transferable-only, so this will be calculated in apply_to
					} else {
						// Unweighted rules
						// Multiply ballots in by surplus fraction
						column.exhausted.votes_out = multiply_surpfrac(column.exhausted.ballots.clone(), &self.surpfrac_numer, &self.surpfrac_denom);
					}
					
					// Round if required
					if let Some(dps) = opts.round_votes {
						column.exhausted.votes_out.floor_mut(dps);
					}
				}
				
				// Sum total votes_out per candidate
				for (candidate, cell) in self.total.cells.iter_mut() {
					cell.votes_out = self.columns.iter().fold(N::new(), |mut acc, col| {
						if let Some(cell) = col.cells.get(candidate) {
							acc += &cell.votes_out;
						}
						acc
					});
				}
				
				// Sum total exhausted votes
				self.total.exhausted.votes_out = self.columns.iter().fold(N::new(), |mut acc, col| { acc += &col.exhausted.votes_out; acc });
			}
		}
	}
	
	/// 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 {
		let mut checksum = N::new();
		
		// Credit transferred votes
		for (candidate, count_card) in state.candidates.iter_mut() {
			if let Some(cell) = self.total.cells.get(candidate) {
				count_card.transfer(&cell.votes_out);
				count_card.ballot_transfers += &cell.ballots;
				checksum += &cell.votes_out;
			}
		}
		
		// Credit exhausted votes
		// If exclusion or not --transferable-only
		if self.surplus.is_none() || !opts.transferable_only {
			// Standard rules
			state.exhausted.transfer(&self.total.exhausted.votes_out);
			state.exhausted.ballot_transfers += &self.total.exhausted.ballots;
			checksum += &self.total.exhausted.votes_out;
		} else {
			// Credit only nontransferable difference
			if self.surpfrac_numer.is_none() {
				// TODO: Is there a purer way of calculating this?
				let difference = self.surplus.as_ref().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;
	}
	
	/// Render table as [Table]
	fn render(&self, opts: &STVOptions) -> Table {
		let mut table = Table::new();
		set_table_format(&mut table);
		
		let show_transfers_per_column = opts.round_subtransfers != RoundSubtransfersMode::SingleStep;
		
		let num_cols;
		if show_transfers_per_column {
			num_cols = self.columns.len() * 3 + 4;
		} else {
			if self.surpfrac.is_none() {
				num_cols = self.columns.len() * 2 + 3;
			} else {
				num_cols = self.columns.len() * 2 + 4;
			}
		}
		
		// ----------
		// Header row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Preference"));
		for column in self.columns.iter() {
			row.push(Cell::new(&format!("Ballots @ {:.dps2$}", column.value_fraction, dps2=max(opts.pp_decimals, 2))).style_spec("cH2"));
			
			if show_transfers_per_column {
				if self.surplus.is_some() {
					row.push(Cell::new(&format!("× {:.dps2$}", self.surpfrac.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2))).style_spec("r"));
				} else {
					row.push(Cell::new("=").style_spec("c"));
				}
			}
		}
		row.push(Cell::new("Total").style_spec("cH2"));
		if self.surpfrac.is_some() {
			row.push(Cell::new(&format!("× {:.dps2$}", self.surpfrac.as_ref().unwrap(), dps2=max(opts.pp_decimals, 2))).style_spec("r"));
		} else if show_transfers_per_column {
			row.push(Cell::new("=").style_spec("c"));
		}
		table.set_titles(Row::new(row));
		
		// --------------
		// Candidate rows
		
		for candidate in self.hopefuls.iter() {
			let mut row = Vec::with_capacity(num_cols);
			row.push(Cell::new(&candidate.name));
			for column in self.columns.iter() {
				if let Some(cell) = column.cells.get(candidate) {
					row.push(Cell::new(&format!("{:.0}", cell.ballots)).style_spec("r"));
					row.push(Cell::new(&format!("{:.dps$}", cell.votes_in, dps=opts.pp_decimals)).style_spec("r"));
					if show_transfers_per_column {
						row.push(Cell::new(&format!("{:.dps$}", cell.votes_out, dps=opts.pp_decimals)).style_spec("r"));
					}
				} else {
					row.push(Cell::new(""));
					row.push(Cell::new(""));
					if show_transfers_per_column {
						row.push(Cell::new(""));
					}
				}
			}
			
			// Totals
			if let Some(cell) = self.total.cells.get(candidate) {
				row.push(Cell::new(&format!("{:.0}", cell.ballots)).style_spec("r"));
				row.push(Cell::new(&format!("{:.dps$}", cell.votes_in, dps=opts.pp_decimals)).style_spec("r"));
				if self.surpfrac.is_some() || show_transfers_per_column {
					row.push(Cell::new(&format!("{:.dps$}", cell.votes_out, dps=opts.pp_decimals)).style_spec("r"));
				}
			} else {
				row.push(Cell::new(""));
				row.push(Cell::new(""));
				if self.surpfrac.is_some() || show_transfers_per_column {
					row.push(Cell::new(""));
				}
			}
			
			table.add_row(Row::new(row));
		}
		
		// -------------
		// Exhausted row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Exhausted"));
		for column in self.columns.iter() {
			if !column.exhausted.ballots.is_zero() {
				row.push(Cell::new(&format!("{:.0}", column.exhausted.ballots)).style_spec("r"));
				row.push(Cell::new(&format!("{:.dps$}", column.exhausted.votes_in, dps=opts.pp_decimals)).style_spec("r"));
				if show_transfers_per_column {
					if column.exhausted.votes_out.is_zero() {
						row.push(Cell::new("-").style_spec("c"));
					} else {
						row.push(Cell::new(&format!("{:.dps$}", column.exhausted.votes_out, dps=opts.pp_decimals)).style_spec("r"));
					}
				}
			} else {
				row.push(Cell::new(""));
				row.push(Cell::new(""));
				if show_transfers_per_column {
					row.push(Cell::new(""));
				}
			}
		}
		
		// Totals
		if !self.total.exhausted.ballots.is_zero() {
			row.push(Cell::new(&format!("{:.0}", self.total.exhausted.ballots)).style_spec("r"));
			row.push(Cell::new(&format!("{:.dps$}", self.total.exhausted.votes_in, dps=opts.pp_decimals)).style_spec("r"));
			if self.surpfrac.is_some() || show_transfers_per_column {
				if self.total.exhausted.votes_out.is_zero() {
					row.push(Cell::new("-").style_spec("c"));
				} else {
					row.push(Cell::new(&format!("{:.dps$}", self.total.exhausted.votes_out, dps=opts.pp_decimals)).style_spec("r"));
				}
			}
		} else {
			row.push(Cell::new(""));
			row.push(Cell::new(""));
			if self.surpfrac.is_some() || show_transfers_per_column {
				row.push(Cell::new(""));
			}
		}
		
		table.add_row(Row::new(row));
		
		// ----------
		// Totals row
		
		let mut row = Vec::with_capacity(num_cols);
		row.push(Cell::new("Total"));
		
		for column in self.columns.iter() {
			row.push(Cell::new(&format!("{:.0}", column.total.ballots)).style_spec("r"));
			row.push(Cell::new(&format!("{:.dps$}", column.total.votes_in, dps=opts.pp_decimals)).style_spec("r"));
			if show_transfers_per_column {
				row.push(Cell::new(&format!("{:.dps$}", column.total.votes_out, dps=opts.pp_decimals)).style_spec("r"));
			}
		}
		
		// Grand total cell
		
		let mut gt_ballots = N::new();
		let mut gt_votes_in = N::new();
		let mut gt_votes_out = N::new();
		
		for candidate in self.hopefuls.iter() {
			if let Some(cell) = self.total.cells.get(candidate) {
				gt_ballots += &cell.ballots;
				gt_votes_in += &cell.votes_in;
				gt_votes_out += &cell.votes_out;
			}
		}
		gt_ballots += &self.total.exhausted.ballots;
		gt_votes_in += &self.total.exhausted.votes_in;
		gt_votes_out += &self.total.exhausted.votes_out;
		
		row.push(Cell::new(&format!("{:.0}", gt_ballots)).style_spec("r"));
		row.push(Cell::new(&format!("{:.dps$}", gt_votes_in, dps=opts.pp_decimals)).style_spec("r"));
		if self.surpfrac.is_some() || show_transfers_per_column {
			row.push(Cell::new(&format!("{:.dps$}", gt_votes_out, dps=opts.pp_decimals)).style_spec("r"));
		}
		
		table.add_row(Row::new(row));
		
		return table;
	}
	
	/// Render table as plain text
	#[cfg(not(target_arch = "wasm32"))]
	pub fn render_text(&self, opts: &STVOptions) -> String {
		return self.render(opts).to_string();
	}
	
	/// Render table as HTML
	#[cfg(target_arch = "wasm32")]
	pub fn render_text(&self, opts: &STVOptions) -> String {
		return self.render(opts).to_html();
	}
	
	// Render table as HTML
	//pub fn render_html(&self, opts: &STVOptions) -> String {
	//	return self.render(opts).to_string();
	//}
}

/// Multiply the specified number by the surplus fraction (if applicable)
fn multiply_surpfrac<N: Number>(mut number: N, surpfrac_numer: &Option<N>, surpfrac_denom: &Option<N>) -> N {
	if let Some(n) = surpfrac_numer {
		number *= n;
	}
	if let Some(n) = surpfrac_denom {
		number /= n;
	}
	return number;
}

/// Column in a [TransferTable]
#[derive(Clone)]
pub struct TransferTableColumn<'e, N: Number> {
	/// Value fraction of ballots counted in this column
	pub value_fraction: N,
	
	/// Number to separate parcels in modes where subtransfers by parcel, etc.
	pub order: usize,
	
	/// Cells in this column
	pub cells: CandidateMap<'e, TransferTableCell<N>>,
	
	/// Exhausted cell
	pub exhausted: TransferTableCell<N>,
	
	/// Totals cell
	pub total: 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(),
				votes_in: N::new(),
				votes_out: N::new(),
			};
			self.cells.insert(candidate, cell);
		}
	}
}

/// Cell in a [TransferTable], representing transfers to one candidate at a particular value
#[derive(Clone)]
pub struct TransferTableCell<N: Number> {
	/// Ballots expressing a next preference for the continuing candidate
	pub ballots: N,
	/// Value of votes when received by the transferring candidate
	pub votes_in: N,
	/// Votes transferred to the continuing candidate
	pub votes_out: N,
}

#[cfg(not(target_arch = "wasm32"))]
fn set_table_format(table: &mut Table) {
	table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
}

#[cfg(target_arch = "wasm32")]
fn set_table_format(_table: &mut Table) {
	// No op
}