use std::rc::Rc;
use itertools::Itertools;

use crate::{ballot::Ballot, header::Header, util::ScoreItem};

#[derive(Debug)]
pub enum Event {
	Lose(ScoreItem),
	Win(Vec<ScoreItem>),
}

#[derive(Debug)]
pub struct Counter {
	pub header: Rc<Header>,
	pub ballots: Vec<Ballot>,
	winners_left: usize,
	enabled: Vec<bool>,
	quota: f64,
}

impl Counter {
	pub fn new<T: std::io::BufRead>(mut csv: quick_csv::Csv<T>, winners: usize) -> Result<Self, Box<dyn std::error::Error>> {
		let header = Rc::new(Header::parse(csv.next().ok_or("csv header missing")??, winners)?);
		let mut ballots = Vec::new();

		if winners > header.candidates.len() {
			return Err("winners can't be smaller than the candidates list".into());
		}

		for row in csv {
			ballots.push(Ballot::parse(row?, &header)?);
		}

		let enabled = vec![true; header.candidates.len()];
		let quota = (ballots.len() as f64) / (header.winners as f64 + 1.0) + 1.0;
		let winners_left = header.winners;

		Ok(Counter {
			header,
			ballots,
			winners_left,
			enabled,
			quota,
		})
	}
	fn count_primaries(&self) -> Vec<ScoreItem> {
		let mut scores = vec![0.0; self.enabled.len()];
		for ballot in self.ballots.iter() {
			if let Some(index) = ballot.get_primary_index(|id| self.enabled[id]) {
				scores[index] += ballot.get_weight();
			}
		}
		scores.into_iter().enumerate().filter_map(|(index, value)| match self.enabled[index] {
			true => Some(ScoreItem::new(index, value)),
			false => None,
		}).collect_vec()
	}
	pub fn get_quota(&self) -> f64 {
		self.quota
	}
}

impl Iterator for Counter {
    type Item = Event;

    fn next(&mut self) -> Option<Self::Item> {
		let mut winners = Vec::<ScoreItem>::with_capacity(self.enabled.len());
		let primaries = self.count_primaries();
		let mut lowest = *primaries.first()?;
		
		for &score in primaries.iter() {
			lowest = lowest.min(score);
			if score.value >= self.quota {
				winners.push(score);
			}
		}

		if primaries.len() == 1 && self.winners_left == 1 {
			self.enabled[lowest.index] = false;
			return Some(Event::Win(vec![lowest]));
		}

		if winners.len() == 0 {
			self.enabled[lowest.index] = false;
			return Some(Event::Lose(lowest));
		}

		winners.sort_by(|a,b| b.cmp(a));

		for ballot in self.ballots.iter_mut() {
			let index = match ballot.get_primary_index(|id| self.enabled[id]) {
				Some(v) => v,
				None => continue,
			};
			if let Some(winner) = winners.iter().filter(|&v| v.index == index).next() {
				let surplus = winner.value - self.quota;
				ballot.apply_weight(surplus / winner.value);
			}
		}

		self.winners_left -= winners.len();
		for winner in winners.iter() {
			self.enabled[winner.index] = false;
		}
		
		Some(Event::Win(winners))
    }
}