wordle/wordle.py

#!/usr/bin/python

import random
import string

from enum import Enum
from collections import defaultdict

WORDLE_LENGTH = 5


class LetterResult(Enum):
    """Is an individual letter in a word in the result?"""
    CORRECT = 'C'
    MOVE = 'M'
    INCORRECT = 'I'
    UNKNOWN = 'U'


class InvalidWord(Exception):
    """Thrown when an input word is not a valid Wordle."""
    pass


_wordle_dictionary = set()


_invalid_words = set(['staia', 'erept', 'boily', 'bolty', 'pindy', 'hyrse', 'hirse', 'retan', 'sisel',
                      'skeel', 'telae', 'genae', 'beata', 'beala', 'macco', 'beele', 'divet', 'vedet',
                      'stail', 'roter', 'fikie', 'muang', 'nurry', 'arara', 'araru', 'araua', 'arawa',
                      'mitra', 'gotra', 'whaly', 'whalm', 'doray', 'korah', 'caroa', 'carlo', 'carot',
                      'carum', 'theet', 'tebet', 'ghole', 'odell'])


def dictionary():
    if not _wordle_dictionary:
        with open('words_alpha.txt') as word_file:
            all_words = word_file.read().split()
        _wordle_dictionary.update(
            filter(lambda w: len(w) == WORDLE_LENGTH and w not in _invalid_words, all_words))

    return _wordle_dictionary


def select_word():
    return random.choice(list(dictionary()))


def grade_word(target, guess):
    # We don't check word length, since all valid words are 5 letters.
    if guess not in dictionary():
        raise InvalidWord(f"{guess} is not a valid Wordle.")

    result = [LetterResult.UNKNOWN] * WORDLE_LENGTH

    # We copy target, and remove letters if we find them. This handles duplicate letters correctly.
    # For example, if the target is "pulls" and we guess "apple", only the first "p" should be a
    # MOVE. The second should be INCORRECT.
    remaining = list(target)

    # Find all correct results first.
    for i in range(WORDLE_LENGTH):
        if target[i] == guess[i]:
            result[i] = LetterResult.CORRECT
            remaining.remove(guess[i])

    for i in range(WORDLE_LENGTH):
        if result[i] == LetterResult.UNKNOWN:
            if guess[i] in remaining:
                # Letter is in the wrong place.
                result[i] = LetterResult.MOVE
                remaining.remove(guess[i])
            else:
                # Letter is not in target.
                result[i] = LetterResult.INCORRECT

    return result


def display_result(result, text):
    opening = {
        LetterResult.CORRECT: '\033[92m',
        LetterResult.MOVE: '\033[33m',
        LetterResult.INCORRECT: '\033[36m',
        LetterResult.UNKNOWN: '',
    }
    closing = '\033[0m'

    return f"{opening[result]}{text}{closing}"


def success_result(result):
    for r in result:
        if r != LetterResult.CORRECT:
            return False

    return True


def analyze_dictionary(wordle_dictionary):
    # Word Position, Letter: Count
    position_matches = [defaultdict(int),
                        defaultdict(int),
                        defaultdict(int),
                        defaultdict(int),
                        defaultdict(int)]

    # Letter: Count
    word_matches = defaultdict(int)

    for word in wordle_dictionary:
        unique_letters = set(word)
        for letter in unique_letters:
            word_matches[letter] += 1

        for index in range(len(word)):
            letter = word[index]
            assert (letter in string.ascii_lowercase)
            position_matches[index][letter] += 1

    return position_matches, word_matches


def score_word(position_matches, word_matches, word):
    result = 0

    # Green Score
    for index in range(len(word)):
        letter = word[index]
        green_score = position_matches[index][letter]
        result += 10 * green_score

    # Yellow Score
    for letter in set(word):
        yellow_score = word_matches[letter]
        result += yellow_score

    return result


def scored_words(wordle_dictionary=dictionary()):
    # word: score
    scored_words = dict()

    position_matches, word_matches = analyze_dictionary(wordle_dictionary)

    for word in wordle_dictionary:
        scored_words[word] = score_word(position_matches, word_matches, word)

    result = list(wordle_dictionary)
    result.sort(key=lambda word: scored_words[word], reverse=True)

    return [(word, scored_words[word]) for word in result]


class Constraints():

    def __init__(self):
        # Array of known letters in result with None for unknown.
        self.correct = [None] * WORDLE_LENGTH
        # Array with per-letter sets of known bad letters.
        self.incorrect = [set() for _ in range(WORDLE_LENGTH)]
        # Dict of letter -> int with minimum required count of letter in word.
        self.count = defaultdict(int)
        # If letter is present in set, value in min is both min and max.
        self.exact = set()

    # List of LetterResult values.
    def update_graded(self, guess, grade):

        guess = guess.lower()

        new_count = defaultdict(int)

        for i, letter in enumerate(guess):
            match grade[i]:
                case LetterResult.CORRECT:
                    self.correct[i] = letter
                    new_count[letter] += 1
                case LetterResult.MOVE:
                    self.incorrect[i].add(letter)
                    new_count[letter] += 1
                case LetterResult.INCORRECT:
                    self.incorrect[i].add(letter)
                    self.exact.add(letter)
                    new_count[letter] += 0  # Ensure letter is in dict

        for letter, c in new_count.items():
            self.count[letter] = max(c, self.count[letter])

    def matches(self, guess):

        counts = defaultdict(int)

        # All known correct letters in place.
        for i, letter in enumerate(guess):
            # If we have a known letter, and don't match...
            if self.correct[i] and self.correct[i] != letter:
                return False

            # If we have a letter that's known to be incorrect in the current location...
            if letter in self.incorrect[i]:
                return False

            counts[letter] += 1

        for letter, c in self.count.items():
            if letter in self.exact:
                if counts[letter] != c:
                    return False
            else:
                if counts[letter] < c:
                    return False

        return True


def filter_dictionary(constraints, wordle_dictionary=dictionary()):
    return set(filter(lambda w: constraints.matches(w), wordle_dictionary))


if __name__ == '__main__':
    constraints = Constraints()
    constraints.update_graded('cares', [LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT])
    constraints.update_graded('monty', [LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT])
    constraints.update_graded('filii', [LetterResult.MOVE, LetterResult.MOVE, LetterResult.INCORRECT, LetterResult.INCORRECT, LetterResult.INCORRECT])

    valid_words = filter_dictionary(constraints)

    print("Top: ", scored_words(valid_words)[:1])