wordle-solve.rs

use std::cmp;
use std::collections::HashMap;
use std::env;
use std::fs;
use std::fs::File;
use std::hash::Hash;
use std::io::Write;
use std::time::SystemTime;

const POWERS_OF_THREE: [u8; 5] = [1, 3, 9, 27, 81];

#[derive(Eq, Hash, PartialEq)]
struct PossibleAnswerSet {
    possible_answers: Vec<u16>,
}

fn score(guess: &Vec<char>, answer: &Vec<char>) -> u8 {
    let mut soln: u8 = 0;
    let mut unpaired_guess_indices = Vec::new();
    let mut unpaired_answer_letters = Vec::new();

    for idx in 0..guess.len() {
        if guess[idx] == answer[idx] {
            soln += 2 * POWERS_OF_THREE[idx];
        } else {
            unpaired_guess_indices.push(idx);
            unpaired_answer_letters.push(answer[idx]);
        }
    }

    for raw_idx in 0..unpaired_guess_indices.len() {
        let guess_idx = unpaired_guess_indices[raw_idx];
        for answer_idx in 0..unpaired_answer_letters.len() {
            if guess[guess_idx] == unpaired_answer_letters[answer_idx] {
                soln += POWERS_OF_THREE[guess_idx];
                unpaired_answer_letters[answer_idx] = unpaired_answer_letters[unpaired_answer_letters.len() - 1];
                unpaired_answer_letters.pop();
                break;
            }
        }
    }

    soln
}

fn format_cmp(raw_comparison: u8) -> String {
    let mut s = String::from("");
    let mut r = raw_comparison;
    for _ in 0..5 {
        match r % 3 {
            0 => s.push('.'),
            1 => s.push('*'),
            2 => s.push('O'),
            _ => panic!(),
        }
        r = r / 3;
    }
    s
}

fn unformat_cmp(formatted_comparison: &Vec<char>) -> u8 {
    let mut r = 0;
    for i in 0..5 {
        if formatted_comparison[i] == '*' {
            r += POWERS_OF_THREE[i];
        } else if formatted_comparison[i] == 'O' {
            r += 2 * POWERS_OF_THREE[i];
        }
    }
    r
}

fn answers_by_result(
    answer_bank: &Vec<Vec<char>>,
    possible_answers: &Vec<u16>,
    guess: &Vec<char>
) -> HashMap<u8, Vec<u16>> {
    let mut h = HashMap::new();

    for answer in possible_answers {
        let answer_str = &answer_bank[*answer as usize];
        let cmp = score(guess, answer_str);
        if !h.contains_key(&cmp) {
            h.insert(cmp, Vec::<u16>::new());
        }
        h.get_mut(&cmp).unwrap().push(*answer as u16);
    }
    h
}

fn search(
    answer_bank: &Vec<Vec<char>>,
    guess_bank: &Vec<Vec<char>>,
    possible_answers: &Vec<u16>,
    stored_strategies: &mut HashMap<PossibleAnswerSet, (f32, String)>,
) -> (f32, String) {
    if possible_answers.len() == 1 {
        return (1.0, String::new());
    }
    if possible_answers.len() == 2 {
        return (1.5, answer_bank[possible_answers[0] as usize].iter().collect());
    }

    let possible_answer_set = PossibleAnswerSet { possible_answers: possible_answers.clone() };
    if stored_strategies.contains_key(&possible_answer_set) {
        return stored_strategies[&possible_answer_set].clone();
    }

    let mut best_result = (1000.0, String::from(""));
    let mut lowest_possibilities = 255;

    for guess in guess_bank {
        let h = answers_by_result(&answer_bank, &possible_answers, guess);
        let most_possibilities = h.keys().map(|k| h[k].len()).max().unwrap();
        lowest_possibilities = cmp::min(lowest_possibilities, most_possibilities);
    }

    for guess in guess_bank {
        let h = answers_by_result(&answer_bank, &possible_answers, guess);

        let most_possibilities = h.keys().map(|k| h[k].len()).max().unwrap();
        if most_possibilities as f32 > 1.4 * (lowest_possibilities as f32) + 2.0 {
            continue;
        }
        if most_possibilities as f32 > 0.5 * (possible_answers.len() as f32) {
            continue;
        }

        let mut sum_weights = 0.0;
        for clue in h.keys() {
            let result_from_clue = search(answer_bank, guess_bank, &h[clue], stored_strategies);
            sum_weights += result_from_clue.0 * h[clue].len() as f32;
        }
        let mut expected_value = sum_weights / possible_answers.len() as f32;
        let mut guess_in_result = false;
        for g in possible_answers {
            if &answer_bank[*g as usize] == guess {
                guess_in_result = true;
                break
            }
        }
        if guess_in_result {
            expected_value -= 1.0 / possible_answers.len() as f32;
        }

        let worst_result = (expected_value, guess.iter().collect());
        if best_result.0 > worst_result.0 {
            best_result = worst_result
        }
    }

    let best_result = (best_result.0 + 1.0, best_result.1);
    let p = PossibleAnswerSet { possible_answers: possible_answers.clone() };
    stored_strategies.insert(p, best_result.clone());
    best_result
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let answers = fs::read_to_string(&args[2]).expect("Reading answers went wrong");
    let answers: Vec<Vec<char>> = answers.split("\n").map(|x| x.trim().chars().collect()).collect();

    let words = fs::read_to_string("words.txt").expect("Something went wrong");
    let words: Vec<Vec<char>> = words.split("\n").map(|x| x.trim().chars().collect()).collect();

    let mut stored_results = vec![vec![0; answers.len()]; words.len()];

    let mut stored_strategies = HashMap::<PossibleAnswerSet, (f32, String)>::new();

    let now = SystemTime::now();
    for answer_idx in 0..answers.len() {
        for word_idx in 0..words.len() {
            stored_results[word_idx][answer_idx] = score(
                &words[word_idx],
                &answers[answer_idx]
            );
        }
    }

    println!("Finished preloading results in {}s", SystemTime::now().duration_since(now).expect("Invalid time").as_secs());
    let now = SystemTime::now();

    search(
        &answers,
        &words,
        &(0..answers.len()).map(|x| x as u16).collect(),
        &mut stored_strategies,
    );

    println!("Finished search in {}s", SystemTime::now().duration_since(now).expect("Invalid time").as_secs());
    let mut fd = File::create(String::from("out_") + &args[2][..]).unwrap();

    for k in stored_strategies.keys() {
        println!("{:?} | {:?}", k.possible_answers.iter().map(|&x| answers[x as usize].iter().collect::<String>()).collect::<Vec<String>>(), stored_strategies[k]);
        writeln!(&mut fd, "{:?} | {:?}", k.possible_answers.iter().map(|&x| answers[x as usize].iter().collect::<String>()).collect::<Vec<String>>(), stored_strategies[k]).unwrap();
    }
}