/*  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::constraints::Constraints;
use crate::election::{CandidateState, CountState, Election, StageKind};
use crate::numbers::{Fixed, GuardedFixed, NativeFloat64, Number, Rational};
use crate::parser::{bin, blt};
use crate::stv::{self, STVOptions};
use crate::ties;

use clap::{AppSettings, Clap};
use itertools::Itertools;

use std::cmp::max;
use std::fs::File;
use std::io::{self, BufRead};
use std::ops;

/// Count a single transferable vote (STV) election
#[derive(Clap)]
#[clap(setting=AppSettings::DeriveDisplayOrder)]
pub struct SubcmdOptions {
	// ----------------
	// -- File input --
	
	/// Path to the BLT file to be counted
	#[clap(help_heading=Some("INPUT"))]
	filename: String,
	
	/// Input is in serialised binary format from "opentally convert"
	#[clap(help_heading=Some("INPUT"), long)]
	bin: bool,
	
	// ----------------------
	// -- Numbers settings --
	
	/// Numbers mode
	#[clap(help_heading=Some("NUMBERS"), short, long, possible_values=&["rational", "fixed", "gfixed", "float64"], default_value="rational", value_name="mode")]
	numbers: String,
	
	/// Decimal places if --numbers fixed
	#[clap(help_heading=Some("NUMBERS"), long, default_value="5", value_name="dps")]
	decimals: usize,
	
	/// Convert ballots with value >1 to multiple ballots of value 1
	#[clap(help_heading=Some("NUMBERS"), long)]
	normalise_ballots: bool,
	
	// -----------------------
	// -- Rounding settings --
	
	/// Round surplus fractions to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, alias="round-tvs", value_name="dps")]
	round_surplus_fractions: Option<usize>,
	
	/// Round ballot values to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, alias="round-weights", value_name="dps")]
	round_values: Option<usize>,
	
	/// Round votes to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_votes: Option<usize>,
	
	/// Round quota to specified decimal places
	#[clap(help_heading=Some("ROUNDING"), long, value_name="dps")]
	round_quota: Option<usize>,
	
	/// (Gregory STV) How to calculate votes to credit to candidates in surplus transfers
	#[clap(help_heading=Some("ROUNDING"), long, possible_values=&["by_value", "per_ballot"], default_value="by_value", value_name="mode")]
	sum_surplus_transfers: String,
	
	/// (Meek STV) Limit for stopping iteration of surplus distribution
	#[clap(help_heading=Some("ROUNDING"), long, default_value="0.001%", value_name="tolerance")]
	meek_surplus_tolerance: String,
	
	// -----------
	// -- Quota --
	
	/// Quota type
	#[clap(help_heading=Some("QUOTA"), short, long, possible_values=&["droop", "hare", "droop_exact", "hare_exact"], default_value="droop")]
	quota: String,
	
	/// Whether to elect candidates on meeting (geq) or strictly exceeding (gt) the quota
	#[clap(help_heading=Some("QUOTA"), short='c', long, possible_values=&["geq", "gt"], default_value="gt", value_name="criterion")]
	quota_criterion: String,
	
	/// Whether to apply a form of progressive quota
	#[clap(help_heading=Some("QUOTA"), long, possible_values=&["static", "ers97", "ers76", "dynamic_by_total", "dynamic_by_active"], default_value="static", value_name="mode")]
	quota_mode: String,
	
	// ------------------
	// -- STV variants --
	
	/// Tie-breaking method
	#[clap(help_heading=Some("STV VARIANTS"), short='t', long, possible_values=&["forwards", "backwards", "random", "prompt"], default_value="prompt", value_name="methods")]
	ties: Vec<String>,
	
	/// Random seed to use with --ties random
	#[clap(help_heading=Some("STV VARIANTS"), long, value_name="seed")]
	random_seed: Option<String>,
	
	/// Method of surplus distributions
	#[clap(help_heading=Some("STV VARIANTS"), short='s', long, possible_values=&["wig", "uig", "eg", "meek", "cincinnati", "hare"], default_value="wig", value_name="method")]
	surplus: String,
	
	/// (Gregory STV) Order to distribute surpluses
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["by_size", "by_order"], default_value="by_size", value_name="order")]
	surplus_order: String,
	
	/// (Gregory STV) Examine only transferable papers during surplus distributions
	#[clap(help_heading=Some("STV VARIANTS"), long)]
	transferable_only: bool,
	
	/// (Gregory STV) Method of exclusions
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["single_stage", "by_value", "by_source", "parcels_by_order", "wright"], default_value="single_stage", value_name="method")]
	exclusion: String,
	
	/// (Meek STV) NZ Meek STV behaviour: Iterate keep values one round before candidate exclusion
	#[clap(help_heading=Some("STV VARIANTS"), long)]
	meek_nz_exclusion: bool,
	
	/// (Cincinnati/Hare) Method of drawing a sample
	#[clap(help_heading=Some("STV VARIANTS"), long, possible_values=&["stratified", "by_order", "nth_ballot"], default_value="stratified", value_name="method")]
	sample: String,
	
	/// (Cincinnati/Hare) Sample-based methods: Check for candidate election after each individual ballot paper transfer
	#[clap(help_heading=Some("STV VARIANTS"), long)]
	sample_per_ballot: bool,
	
	// -------------------------
	// -- Count optimisations --
	
	/// Continue count even if continuing candidates fill all remaining vacancies
	#[clap(help_heading=Some("COUNT OPTIMISATIONS"), long)]
	no_early_bulk_elect: bool,
	
	/// Use bulk exclusion
	#[clap(help_heading=Some("COUNT OPTIMISATIONS"), long)]
	bulk_exclude: bool,
	
	/// Defer surplus distributions if possible
	#[clap(help_heading=Some("COUNT OPTIMISATIONS"), long)]
	defer_surpluses: bool,
	
	/// Elect candidates only when their surpluses are distributed; (Meek STV) Wait for keep values to converge before electing candidates
	#[clap(help_heading=Some("COUNT OPTIMISATIONS"), long)]
	no_immediate_elect: bool,
	
	/// On exclusion, exclude any candidate with fewer than this many votes
	#[clap(help_heading=Some("COUNT OPTIMISATIONS"), long, default_value="0", value_name="votes")]
	min_threshold: String,
	
	// -----------------
	// -- Constraints --
	
	/// Path to a CON file specifying constraints
	#[clap(help_heading=Some("CONSTRAINTS"), long)]
	constraints: Option<String>,
	
	/// Mode of handling constraints
	#[clap(help_heading=Some("CONSTRAINTS"), long, possible_values=&["guard_doom"], default_value="guard_doom")]
	constraint_mode: String,
	
	// ---------------------
	// -- Output settings --
	
	#[clap(help_heading=Some("OUTPUT"), short, long, possible_values=&["text", "csv"], default_value="text")]
	output: String,
	
	/// Hide excluded candidates from results report
	#[clap(help_heading=Some("OUTPUT"), long)]
	hide_excluded: bool,
	
	/// Sort candidates by votes in results report
	#[clap(help_heading=Some("OUTPUT"), long)]
	sort_votes: bool,
	
	/// Print votes to specified decimal places in results report
	#[clap(help_heading=Some("OUTPUT"), long, default_value="2", value_name="dps")]
	pp_decimals: usize,
}

/// Entrypoint for subcommand
pub fn main(cmd_opts: SubcmdOptions) -> Result<(), i32> {
	// Read and count election according to --numbers
	if cmd_opts.numbers == "rational" {
		let mut election = election_from_file(&cmd_opts.filename, cmd_opts.bin)?;
		maybe_load_constraints(&mut election, &cmd_opts.constraints);
		
		// Must specify ::<N> here and in a few other places because ndarray causes E0275 otherwise
		count_election::<Rational>(election, cmd_opts)?;
	} else if cmd_opts.numbers == "float64" {
		let mut election = election_from_file(&cmd_opts.filename, cmd_opts.bin)?;
		maybe_load_constraints(&mut election, &cmd_opts.constraints);
		count_election::<NativeFloat64>(election, cmd_opts)?;
	} else if cmd_opts.numbers == "fixed" {
		Fixed::set_dps(cmd_opts.decimals);
		
		let mut election = election_from_file(&cmd_opts.filename, cmd_opts.bin)?;
		maybe_load_constraints(&mut election, &cmd_opts.constraints);
		count_election::<Fixed>(election, cmd_opts)?;
	} else if cmd_opts.numbers == "gfixed" {
		GuardedFixed::set_dps(cmd_opts.decimals);
		
		let mut election = election_from_file(&cmd_opts.filename, cmd_opts.bin)?;
		maybe_load_constraints(&mut election, &cmd_opts.constraints);
		count_election::<GuardedFixed>(election, cmd_opts)?;
	}
	
	return Ok(());
}

fn election_from_file<N: Number>(path: &str, bin: bool) -> Result<Election<N>, i32> {
	if bin {
		// BIN format
		return Ok(bin::parse_path(path));
	} else {
		// BLT format
		match blt::parse_path(path) {
			Ok(e) => return Ok(e),
			Err(err) => {
				println!("Syntax Error: {}", err);
				return Err(1);
			}
		}
	}
}

fn maybe_load_constraints<N: Number>(election: &mut Election<N>, constraints: &Option<String>) {
	if let Some(c) = constraints {
		let file = File::open(c).expect("IO Error");
		let lines = io::BufReader::new(file).lines();
		election.constraints = Some(Constraints::from_con(lines.map(|r| r.expect("IO Error").to_string()).into_iter()));
	}
}

fn count_election<N: Number>(election: Election<N>, cmd_opts: SubcmdOptions) -> Result<(), i32>
where
	for<'r> &'r N: ops::Add<&'r N, Output=N>,
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	// Copy applicable options
	let opts = STVOptions::new(
		cmd_opts.round_surplus_fractions,
		cmd_opts.round_values,
		cmd_opts.round_votes,
		cmd_opts.round_quota,
		cmd_opts.sum_surplus_transfers.into(),
		cmd_opts.meek_surplus_tolerance.into(),
		cmd_opts.normalise_ballots,
		cmd_opts.quota.into(),
		cmd_opts.quota_criterion.into(),
		cmd_opts.quota_mode.into(),
		ties::from_strs(cmd_opts.ties, cmd_opts.random_seed),
		cmd_opts.surplus.into(),
		cmd_opts.surplus_order.into(),
		cmd_opts.transferable_only,
		cmd_opts.exclusion.into(),
		cmd_opts.meek_nz_exclusion,
		cmd_opts.sample.into(),
		cmd_opts.sample_per_ballot,
		!cmd_opts.no_early_bulk_elect,
		cmd_opts.bulk_exclude,
		cmd_opts.defer_surpluses,
		!cmd_opts.no_immediate_elect,
		cmd_opts.min_threshold,
		cmd_opts.constraints,
		cmd_opts.constraint_mode.into(),
		cmd_opts.hide_excluded,
		cmd_opts.sort_votes,
		cmd_opts.pp_decimals,
	);
	
	// Validate options
	match opts.validate() {
		Ok(_) => {}
		Err(err) => {
			println!("Error: {}", err.describe());
			return Err(1);
		}
	}
	
	match cmd_opts.output.as_str() {
		"text" => { return count_election_text(election, &cmd_opts.filename, opts); }
		"csv" => { return count_election_csv(election, opts); }
		_ => unreachable!()
	}
}

// ---------------
// CLI text report

fn count_election_text<N: Number>(mut election: Election<N>, filename: &str, opts: STVOptions) -> Result<(), i32>
where
	for<'r> &'r N: ops::Add<&'r N, Output=N>,
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	// Describe count
	let total_ballots = election.ballots.iter().fold(N::zero(), |acc, b| { acc + &b.orig_value });
	print!("Count computed by OpenTally (revision {}). Read {:.0} ballots from \"{}\" for election \"{}\". There are {} candidates for {} vacancies. ", crate::VERSION, total_ballots, filename, election.name, election.candidates.len(), election.seats);
	let opts_str = opts.describe::<N>();
	if opts_str.len() > 0 {
		println!("Counting using options \"{}\".", opts_str);
	} else {
		println!("Counting using default options.");
	}
	println!();
	
	stv::preprocess_election(&mut election, &opts);
	
	// Initialise count state
	let mut state = CountState::new(&election);
	
	// Distribute first preferences
	match stv::count_init(&mut state, &opts) {
		Ok(_) => {}
		Err(err) => {
			println!("Error: {}", err.describe());
			return Err(1);
		}
	}
	
	let mut stage_num = 1;
	print_stage(stage_num, &state, &opts);
	
	loop {
		match stv::count_one_stage(&mut state, &opts) {
			Ok(is_done) => {
				if is_done {
					break;
				}
			}
			Err(err) => {
				println!("Error: {}", err.describe());
				return Err(1);
			}
		}
		
		stage_num += 1;
		print_stage(stage_num, &state, &opts);
	}
	
	println!("Count complete. The winning candidates are, in order of election:");
	
	let mut winners = Vec::new();
	for (candidate, count_card) in state.candidates.iter() {
		if count_card.state == CandidateState::Elected {
			winners.push((candidate, count_card));
		}
	}
	winners.sort_unstable_by(|a, b| a.1.order_elected.cmp(&b.1.order_elected));
	
	for (i, (winner, count_card)) in winners.into_iter().enumerate() {
		if let Some(kv) = &count_card.keep_value {
			println!("{}. {} (kv = {:.dps2$})", i + 1, winner.name, kv, dps2=max(opts.pp_decimals, 2));
		} else {
			println!("{}. {}", i + 1, winner.name);
		}
	}
	
	return Ok(());
}

fn print_stage<N: Number>(stage_num: u32, state: &CountState<N>, opts: &STVOptions) {
	// Print stage details
	println!("{}. {}", stage_num, state.title);
	println!("{}", state.logger.render().join(" "));
	
	// Print candidates
	print!("{}", state.describe_candidates(opts));
	
	// Print summary rows
	print!("{}", state.describe_summary(opts));
	
	println!("");
}

// ----------------------------------
// Wichmann/eSTV/ERS-style CSV report

fn count_election_csv<N: Number>(mut election: Election<N>, opts: STVOptions) -> Result<(), i32>
where
	for<'r> &'r N: ops::Add<&'r N, Output=N>,
	for<'r> &'r N: ops::Sub<&'r N, Output=N>,
	for<'r> &'r N: ops::Mul<&'r N, Output=N>,
	for<'r> &'r N: ops::Div<&'r N, Output=N>,
	for<'r> &'r N: ops::Neg<Output=N>
{
	// Header rows
	let total_ballots = election.ballots.iter().fold(N::zero(), |acc, b| { acc + &b.orig_value });
	
	// eSTV does not consistently quote records, so we won't use a CSV library here
	println!(r#""Election for","{}""#, election.name);
	println!(r#""Date","  /  /    ""#);
	println!(r#""Number to be elected",{}"#, election.seats);
	
	stv::preprocess_election(&mut election, &opts);
	
	// Initialise count state
	let mut state = CountState::new(&election);
	
	let mut stage_results = vec![Vec::new(); election.candidates.len() + 5];
	
	// -----------
	// First stage
	
	// Distribute first preferences
	match stv::count_init(&mut state, &opts) {
		Ok(_) => {}
		Err(err) => {
			println!("Error: {}", err.describe());
			return Err(1);
		}
	}
	
	// Subtract this from progressive NTs
	// FIXME: May fail to round correctly with minivoters
	let invalid_votes = state.exhausted.votes.clone();
	let valid_votes = total_ballots - &invalid_votes;
	
	// Stage number row
	stage_results[0].push(String::new());
	stage_results[0].push(String::new());
	
	// Stage kind row
	stage_results[1].push(String::new());
	stage_results[1].push(String::from(r#""First""#));
	
	// Stage title row
	stage_results[2].push(String::from(r#""Candidates""#));
	stage_results[2].push(String::from(r#""Preferences""#));
	
	for (i, candidate) in election.candidates.iter().enumerate() {
		let count_card = &state.candidates[candidate];
		stage_results[3 + i].push(format!(r#""{}""#, candidate.name));
		stage_results[3 + i].push(format!(r#"{:.0}"#, count_card.votes)); // FIXME: May fail to round correctly with minivoters
	}
	
	stage_results[3 + election.candidates.len()].push(String::from(r#""Non-transferable""#));
	stage_results[3 + election.candidates.len()].push(String::new()); // FIXME: May fail to round correctly with minivoters
	
	stage_results[4 + election.candidates.len()].push(String::from(r#""Totals""#));
	stage_results[4 + election.candidates.len()].push(format!(r#"{:.0}"#, valid_votes));
	
	//let mut orig_states = HashMap::new();
	//for (candidate, count_card) in state.candidates.iter() {
	//	orig_states.insert(*candidate, count_card.state);
	//}
	
	// -----------------
	// Subsequent stages
	
	let mut stage_num: u32 = 1;
	loop {
		match stv::count_one_stage(&mut state, &opts) {
			Ok(is_done) => {
				if is_done {
					break;
				}
			}
			Err(err) => {
				println!("Error: {}", err.describe());
				return Err(1);
			}
		}
		
		stage_num += 1;
		
		// Stage number row
		stage_results[0].push(String::from(r#""Stage""#));
		stage_results[0].push(format!(r#"{}"#, stage_num));
		
		// Stage kind row
		stage_results[1].push(format!(r#""{}""#, state.title.kind_as_string()));
		stage_results[1].push(String::new());
		
		// Stage title row
		match &state.title {
			StageKind::FirstPreferences => unreachable!(),
			StageKind::SurplusOf(candidate) => {
				stage_results[2].push(format!(r#""{}""#, candidate.name));
			}
			StageKind::ExclusionOf(candidates) => {
				stage_results[2].push(format!(r#""{}""#, candidates.iter().map(|c| &c.name).sorted().join("+")));
			}
			StageKind::SurplusesDistributed => todo!(),
			StageKind::BulkElection => {
				//let mut elected_candidates = Vec::new();
				//for candidate in election.candidates.iter() {
				//	if state.candidates[candidate].state == CandidateState::Hopeful && orig_states[candidate].state != CandidateState::Hopeful {
				//		elected_candidates.push(candidate);
				//	}
				//}
				stage_results[2].push(String::from(r#""Bulk election""#));
			}
		}
		stage_results[2].push(String::from(r#""#));
		
		for (i, candidate) in election.candidates.iter().enumerate() {
			let count_card = &state.candidates[candidate];
			
			if count_card.transfers.is_zero() {
				stage_results[3 + i].push(String::new());
			} else if count_card.transfers > N::zero() {
				stage_results[3 + i].push(format!(r#"+{:.dps$}"#, count_card.transfers, dps=opts.pp_decimals));
			} else {
				stage_results[3 + i].push(format!(r#"{:.dps$}"#, count_card.transfers, dps=opts.pp_decimals));
			}
			
			if count_card.votes.is_zero() {
				stage_results[3 + i].push(String::from(r#""-""#));
			} else {
				stage_results[3 + i].push(format!(r#"{:.dps$}"#, count_card.votes, dps=opts.pp_decimals));
			}
		}
		
		// Nontransferable
		let nt_transfers = state.exhausted.transfers.clone() + &state.loss_fraction.transfers;
		if nt_transfers.is_zero() {
			stage_results[3 + election.candidates.len()].push(String::new());
		} else if nt_transfers > N::zero() {
			stage_results[3 + election.candidates.len()].push(format!(r#"+{:.dps$}"#, nt_transfers, dps=opts.pp_decimals));
		} else {
			stage_results[3 + election.candidates.len()].push(format!(r#"{:.dps$}"#, nt_transfers, dps=opts.pp_decimals));
		}
		stage_results[3 + election.candidates.len()].push(format!(r#"{:.dps$}"#, &state.exhausted.votes + &state.loss_fraction.votes - &invalid_votes, dps=opts.pp_decimals));
		
		// Totals
		stage_results[4 + election.candidates.len()].push(String::new());
		stage_results[4 + election.candidates.len()].push(format!(r#"{:.dps$}"#, valid_votes, dps=opts.pp_decimals));
		
		//for (candidate, count_card) in state.candidates.iter() {
		//	orig_states.insert(*candidate, count_card.state);
		//}
	}
	
	// ----------------
	// Candidate states
	
	stage_results[3 + election.candidates.len()].push(String::new()); // Nontransferable row
	
	for (i, candidate) in election.candidates.iter().enumerate() {
		let count_card = &state.candidates[candidate];
		if count_card.state == CandidateState::Elected {
			stage_results[3 + i].push(String::from(r#""Elected""#));
		} else {
			stage_results[3 + i].push(String::new());
		}
	}
	
	// --------------------
	// Output stages to CSV
	
	println!(r#""Valid votes",{:.0}"#, valid_votes);
	println!(r#""Invalid votes",{:.0}"#, invalid_votes);
	println!(r#""Quota",{:.dps$}"#, state.quota.as_ref().unwrap(), dps=opts.pp_decimals);
	println!(r#""OpenTally","{}""#, crate::VERSION);
	println!(r#""Election rules","{}""#, opts.describe::<N>());
	
	for row in stage_results {
		println!("{}", row.join(","));
	}
	
	return Ok(());
}
