wordle-solver-rs/src/lib/solver.rs

use std::sync::{
    atomic::{AtomicUsize, Ordering},
    mpsc, Arc,
};

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};

#[derive(Clone)]
pub struct Guess {
    pub pattern: Pattern,
    pub information_gained: f64,
}

#[derive(Clone)]
pub struct Solver {
    valid_words: WordList,
    possible_solutions: WordList,
    guesses: Vec<Guess>,
}

pub struct SolverJob {
    receiver: mpsc::Receiver<Result<Vec<WordStats>, String>>,
}

impl SolverJob {
    pub fn wait(self) -> Result<Vec<WordStats>, String> {
        self.receiver
            .recv()
            .map_err(|_| "calculation task cancelled".to_string())?
    }
}

#[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![],
        }
    }

    pub fn new(valid_words: WordList, possible_solutions: WordList) -> Solver {
        Solver {
            valid_words,
            possible_solutions,
            guesses: 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("Applying this guess would empty the solution space!".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_async<F>(&self, progress: F) -> SolverJob
    where
        F: Fn(usize, usize) + Send + Sync + 'static,
    {
        let solver_locked = (*self).clone();

        let (tx, rc) = mpsc::channel();

        let completed = Arc::new(AtomicUsize::new(0));

        rayon::spawn_fifo(move || {
            let mut result = solver_locked
                .valid_words
                .words()
                .par_bridge()
                .map(|w| {
                    let r = solver_locked.evaluate_word(w);
                    let completed = completed.fetch_add(1, Ordering::Relaxed) + 1;
                    progress(completed, solver_locked.valid_words.len());
                    r
                })
                .collect::<Result<Vec<WordStats>, String>>();

            if let Ok(result) = &mut result {
                result.sort_by(|a, b| {
                    a.expected_score_after_guess
                        .total_cmp(&b.expected_score_after_guess)
                });
            }

            let _ = tx.send(result);
        });

        SolverJob { receiver: rc }
    }

    pub fn evaluate_all_words_cb<F>(&self, progress: F) -> Result<Vec<WordStats>, String>
    where
        F: Fn(usize, usize) + Send + Sync + 'static,
    {
        let mut completed = 0;
        let mut result = self
            .valid_words
            .words()
            .map(|w| {
                let r = self.evaluate_word(w);
                completed += 1;
                progress(completed, self.valid_words.len());
                r
            })
            .collect::<Result<Vec<WordStats>, String>>();

        if let Ok(result) = &mut result {
            result.sort_by(|a, b| {
                a.expected_score_after_guess
                    .total_cmp(&b.expected_score_after_guess)
            });
        }
        result
    }

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

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

        Ok(best_words)
    }

    pub fn best_word_format(best_words: Vec<WordStats>, 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 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 guesses(&self) -> &Vec<Guess> {
        &self.guesses
    }

    pub fn stats(&self) -> (u64, f64) {
        (
            self.possible_solutions.len() as u64,
            self.possible_solutions.equal_likelieness_entropy(),
        )
    }

    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
        )
    }
}