use crate::pattern::Pattern;
use indicatif::ParallelProgressIterator;
use rayon::iter::{ParallelBridge, ParallelIterator};

use crate::word_list::WordList;

use derive_more::Display;
use tabular::{Row, Table};

struct Guess {
    pattern: Pattern,
    information_gained: f64,
}

pub struct Solver {
    valid_words: WordList,
    possible_solutions: WordList,
    guesses: Vec<Guess>,
    best_words: Vec<WordStats>,
}

#[derive(Display, Clone)]
#[display(
    "WordStats({}, p={:.02}%, E[I]={:.02}, E[s]={:.02})",
    word,
    solution_probability*100.0,
    expected_information_gained,
    expected_score_after_guess
)]
pub struct WordStats {
    pub word: String,
    pub solution_probability: f64,
    pub expected_information_gained: f64,
    pub expected_score_after_guess: f64,
}

impl Solver {
    pub fn from_single_word_list(word_list: WordList) -> Solver {
        Solver {
            valid_words: word_list.clone(),
            possible_solutions: word_list,
            guesses: vec![],
            best_words: vec![],
        }
    }

    pub fn new(valid_words: WordList, possible_solutions: WordList) -> Solver {
        Solver {
            valid_words,
            possible_solutions,
            guesses: vec![],
            best_words: vec![],
        }
    }

    pub fn apply_guess(&mut self, guess: Pattern) -> Result<(), String> {
        let new_possible_solutions = self.possible_solutions.apply_guess(&guess);
        let information_gained = self.possible_solutions.equal_likelieness_entropy()
            - new_possible_solutions.equal_likelieness_entropy();

        if new_possible_solutions.len() == 0 {
            return Err("Solution space is empty".to_string());
        }

        self.possible_solutions = new_possible_solutions;

        self.guesses.push(Guess {
            pattern: guess,
            information_gained,
        });

        Ok(())
    }

    fn estimate_guesses(&self, expected_information_gained: f64) -> f64 {
        if self.possible_solutions.len() == 1 {
            1.0
        } else {
            1.0 + self.possible_solutions.equal_likelieness_entropy() - expected_information_gained
        }
    }

    pub fn evaluate_word(&self, word: &str) -> Result<WordStats, String> {
        let solution_probability = self.possible_solutions.word_probability(word);

        let expected_information_gained = self.possible_solutions.entropy_if_guessed(word)?;

        let score = (self.guesses.len() + 1) as f64;

        let expected_score_after_guess = score * solution_probability
            + (1.0 - solution_probability)
                * (score + self.estimate_guesses(expected_information_gained));

        Ok(WordStats {
            word: word.to_string(),
            solution_probability,
            expected_information_gained,
            expected_score_after_guess,
        })
    }

    pub fn evaluate_all_words(&mut self) -> Result<Vec<WordStats>, String> {
        self.best_words = self
            .valid_words
            .words()
            .par_bridge()
            .progress_count(self.valid_words.len() as u64)
            .map(|w| self.evaluate_word(w))
            .collect::<Result<_, String>>()?;

        self.best_words.sort_by(|a, b| {
            a.expected_score_after_guess
                .total_cmp(&b.expected_score_after_guess)
        });

        Ok(self.best_words.clone())
    }

    pub fn best_word_format(&self, n: usize) -> String {
        let mut table = Table::new("#{:<} {:>} | {:<} {:<} {:<}");

        table.add_row(
            Row::new()
                .with_cell("Best")
                .with_cell("Word")
                .with_cell("p[solution]")
                .with_cell("E[I]")
                .with_cell("E[score]"),
        );
        for (i, ws) in self.best_words.iter().take(n).enumerate() {
            table.add_row(
                Row::new()
                    .with_cell(i + 1)
                    .with_cell(&ws.word)
                    .with_cell(ws.solution_probability)
                    .with_cell(ws.expected_information_gained)
                    .with_cell(ws.expected_score_after_guess),
            );
        }

        table.to_string()
    }

    pub fn tabular_format(&self) -> String {
        let mut table = Table::new("{:<} | {:<}");

        table.add_row(Row::new().with_cell("Guess").with_cell("Info-Gained"));

        for g in &self.guesses {
            table.add_row(
                Row::new()
                    .with_ansi_cell(g.pattern.ansi_format())
                    .with_cell(format!("{:.04}", g.information_gained)),
            );
        }

        format!(
            "Solution Space: {}, Uncertainty: {:.02} bits\n{}",
            self.possible_solutions.len(),
            self.possible_solutions.equal_likelieness_entropy(),
            table
        )
    }
}