use std::{sync::{Arc, Mutex}, thread};
use itertools::Itertools;

use crate::{ballot::Ballot, csv, header::Header, stats::Stats, util::{self, ScoreItem}};

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

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

const CHUNK_SIZE: usize = 1024;

impl Counter {
	pub fn new(mut csv: csv::reader::CsvReader, winners: usize) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
		let header = Arc::new(Header::parse(csv.next().ok_or("csv header missing")?, winners)?);
		let ballots = Mutex::new(Vec::new());
		let stats = Mutex::new(Stats::new());
		let csv = Mutex::new(csv);

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

		util::thread::spawn_all_with_result(|| -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
			let mut l_stats = Stats::new();
			let mut l_ballots = Vec::new();
			loop {
				let rows = csv.lock().unwrap().by_ref().take(CHUNK_SIZE).collect_vec();

				if rows.len() == 0 {
					break;
				}

				for row in rows {
					l_ballots.push(Ballot::parse(row, &header, &mut l_stats).unwrap());
				}

				ballots.lock().unwrap().append(&mut l_ballots);
			}
			{stats.lock().unwrap().add(&l_stats)};
			Ok(())
		})?;

		let ballots = ballots.into_inner().unwrap();
		let stats = stats.into_inner().unwrap();

		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,
			stats,
			winners_left,
			enabled,
			quota,
		})
	}
	fn count_primaries(&self) -> Vec<ScoreItem> {
		let mut scores = Mutex::new(vec![0.0; self.enabled.len()]);
		let ballot_chunks = Mutex::new(self.ballots.chunks(CHUNK_SIZE));

		util::thread::spawn_all(|| {
			let mut l_scores = vec![0.0; self.enabled.len()];
			for ballot in util::mutex::chunk_iter(&ballot_chunks) {
				ballot.count_primary(&mut l_scores, &self.enabled);
			}
			for (dst, src) in scores.lock().unwrap().iter_mut().zip(l_scores.into_iter()) {
				*dst += src;
			}
		});

		scores.get_mut().unwrap().iter().copied().enumerate().filter_map(|(index, value)| match self.enabled[index] {
			true => Some(ScoreItem::new(index, value)),
			false => None,
		}).collect_vec()
	}
	fn apply_weights(&mut self, winners: &[ScoreItem]) {
		let ballot_chunks = Mutex::new(self.ballots.chunks_mut(CHUNK_SIZE));

		util::thread::spawn_all(|| {
			for ballot in util::mutex::chunk_iter_mut(&ballot_chunks) {
				ballot.apply_winners(&self.enabled, winners, self.quota);
			}
		});
	}
	#[inline]
	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));
		self.apply_weights(&winners);
		self.winners_left -= winners.len();

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