#!/usr/bin/python

from collections import defaultdict

import dictionary
from game import LetterResult


class Constraints:
    def __init__(self):
        # Array of known letters in result with None for unknown.
        self.correct = [None] * dictionary.WORDLE_LENGTH
        # Array with per-letter sets of known bad letters.
        self.incorrect = [set() for _ in range(dictionary.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_words(self, words=dictionary.words()):
        return set(filter(lambda w: self.matches(w), words))