"""RuleSet and associated classes

Classes:
    Rule - rule for a letter of the alphabet
    RuleSet - set of rules for this Wordle game
    WordScore - scoring for a single word

This module does the heavy lifting for the wordlesolve package.
As clues are uncovered it adds rules to the RuleSet e.g.
    - there is no 'P' in the solution
    - there is at least one 'A' and it can't be the first letter in the word
    - the first letter is 'C'

Each rule is represented by an instance of the Rule class.
Rule is a dataclass with attributes:
    - count: number of confirmed occurrences of this letter in the solution
    - count_op: whether count is the exact or minimum number of occurrences
    - confirmed: positions where this letter is confirmed as present
    - excluded: positions where this letter is confirmed as not present

The RuleSet class is a dict subclass where each key is a letter
and the value is the Rule corresponding to that letter.

Rules are generated by adding guess/outcome pairs to the RuleSet.
An outcome is a five-letter string where each string represents
a letter from the guess:
    - a letter not present (grey) is represented by a 0
    - a letter present but not in the correct position (yellow) is represented by a 1
    - a letter present and in the correct position (green) is represeted by a 2

E.g. guess 'BRAKE' and outcome '01200' will generate 5 rules:
{
    'B': Rule(count=0, count_op='eq', confirmed=set(), excluded=set()),
    'R': Rule(count=1, count_op='gte', confirmed=set(), excluded={1}),
    'A': Rule(count=1, count_op='gte', confirmed={2}, excluded=set()),
    'K': Rule(count=0, count_op='eq', confirmed=set(), excluded=set()),
    'E': Rule(count=0, count_op='eq', confirmed=set(), excluded=set()),
}

RuleSet methods:
    .add_outcome(guess, outcome): add an outcome to this RuleSet
    .get_alphabet_status(): returns a dict showing the status of each letter of the alphabet.
    .is_match(word): check whether a word matches this RuleSet
    .filter_matches(words): filter a list of words to retain only RuleSet matches
    .score_words(words, matches): score a list of words according to this RuleSet

RuleSet.score_words(words, matches) returns an instance of the WordScore class.
WordScore is a subclass of NamedTuple with attributes:
    - word: the word being scored
    - score: the score for this word
    - frequency: word frequency (used to tiebreak between words with equal score)

"""

from collections import Counter
from dataclasses import dataclass, field
from operator import itemgetter
import string
from typing import NamedTuple

from wordfreq import zipf_frequency  # type: ignore


@dataclass
class Rule:
    """Rule for a single letter

    Attributes:
        count:
            number of confirmed instances of this letter in the solution

        count_op:
            identifies the number in count as either:
                'gte': the letter appears at least this number of times
                'eq': the letter appears exactly this number of times

        confirmed:
            positions where this letter is confirmed as present

        excluded:
            positions where this letter is confirmed as not present

    E.g. guess BRAKE and outcome 01200 will generate 5 rules:
        B: count=0, count_op='eq', confirmed=set(), excluded=set()
        R: count=1, count_op='gte', confirmed=set(), excluded={1}
        A: count=1, count_op='gte', confirmed={2}, excluded=set()
        K: count=0, count_op='eq', confirmed=set(), excluded=set()
        E: count=0, count_op='eq', confirmed=set(), excluded=set()

    """

    count: int = 0
    count_op: str = "gte"
    confirmed: set[int] = field(default_factory=set)
    excluded: set[int] = field(default_factory=set)


class WordScore(NamedTuple):
    """Score for a single word"""

    word: str  # word being scored
    score: int  # total score
    frequency: int = 0  # word frequency


class RuleSet(dict):
    """Set of Rules

    A dict subclass where each key is a letter and the value
    is the Rule corresponding to that letter
    """

    def add_outcome(self, guess: str, outcome: str):
        """Adds an outcome to the ruleset

        Args:
            guess: five-letter word that has been guessed
            outcome: five-character outcome string
                     characters correspond to letters in guess
                     '0' = letter not present
                     '1' = letter present but not in correct position
                     '2' = letter present and in correct position
        """

        # Get count of letters
        counts = {letter: 0 for letter in set(guess)}
        for letter, result in zip(guess, outcome):
            if result in ("1", "2"):
                counts[letter] += 1

        # Add to ruleset
        for letter, count in counts.items():

            # Create rule if not already there
            if letter not in self:
                self[letter] = Rule()

            # Increase letter count if applicable
            self[letter].count = max(self[letter].count, count)

        # Add positions to ruleset
        for letter, result, position in zip(guess, outcome, range(5)):

            # 0 means we know the maximum occurrences there can be of this letter
            if result == "0":
                self[letter].count_op = "eq"

            # 1 means the letter cannot be in this position
            elif result == "1":
                self[letter].excluded.add(position)

            # 2 means the letter must be in this position
            else:
                self[letter].confirmed.add(position)

    def is_match(self, word: str) -> bool:
        """Checks whether a word matches this rule set

        Args:
            word: word to check for match

        Returns:
            True if the word matches the RuleSet
            False otherwise
        """

        # Check each letter
        for position, letter in enumerate(word):

            # No rule for this letter = passes
            if letter in self:

                rule = self[letter]

                # Letter not in solution
                if rule.count == 0:
                    return False

                # Position excluded
                if position in rule.excluded:
                    return False

                # Matches confirmed position
                for pos in rule.confirmed:
                    if word[pos] != letter:
                        return False

        # Check all required letters are present
        for letter, rule in self.items():
            if rule.count_op == "gte" and word.count(letter) < rule.count:
                return False
            if rule.count_op == "eq" and word.count(letter) != rule.count:
                return False

        return True

    def filter_matches(self, words: list[str]) -> list[str]:
        """Filters a list of words to retain only ruleset matches

        Args:
            words: list of words to filter

        Returns:
            list of those words matching this ruleset
        """

        return [word for word in words if self.is_match(word)]

    def score_words(self, words: list[str], matches: list[str]) -> list[WordScore]:
        """
        Score a list of words according to this RuleSet

        Args:
            words: words to score
            matches: matches to current RuleSet

        Returns:
            list of WordScores ordered from highest to lowest score
        """

        # Get frequency and position scores from the matched words
        fscores = self._get_frequency_scores(matches)
        pscores = self._get_position_scores(matches)

        # If we are down to 1 or 2 matches just score those
        # otherwise score the whole list
        words_to_score = matches if len(matches) < 3 else words

        # Get word scores
        word_scores = _score_words(words_to_score, fscores, pscores)

        # Order by score
        word_scores.sort(key=itemgetter(1), reverse=True)

        # Word frequency breaks tie for words as far as fifth position
        previous = 0
        for index, word_score in enumerate(word_scores):

            # Loop at least to the fifth word
            # and then until the score changes
            if index > 4 and word_score.score != previous:
                break

            # Add word frequency to the score
            new_score = word_score._replace(
                frequency=zipf_frequency(word_score.word, "en")
            )
            word_scores[index] = new_score
            previous = word_score.score

        # Order again by word frequency and then by score
        word_scores.sort(key=itemgetter(2), reverse=True)
        word_scores.sort(key=itemgetter(1), reverse=True)

        return word_scores

    def get_alphabet_status(self) -> dict[str, int]:
        """Returns a dict showing the status of each letter of the alphabet.

        Dict keys are uppercase letters, values indicate the letter's status
        Possible values are:
            -1: letter status unknown
             0: letter is known not to be in the solution
             1: letter is known to be in the solution but its position is not known
             2: letter is known to be in the solution and its position is also known
        """
        return {
            letter: self._get_letter_status(letter) for letter in string.ascii_uppercase
        }

    def _get_letter_status(self, letter) -> int:
        """Returns the status of a single letter.

        Possible values are:
            -1: letter status unknown
             0: letter is known not to be in the solution
             1: letter is known to be in the solution but its position is not known
             2: letter is known to be in the solution and its position is also known
        """

        # No rule = no information about this letter
        if letter not in self:
            return -1

        # Letter not in solution
        if self[letter].count == 0:
            return 0

        # Letter in correct position
        if self[letter].confirmed:
            return 2

        # Letter not in correct position
        return 1

    def _get_frequency_mask(self) -> dict[str, list[int]]:
        """Gets the mask for frequency scores

        The mask indicates which letter occurrences are of relevance when scoring a word.

        Rules are;
        - if a letter's frequency is known, it will not be scored
        - if a letter's minimum frequency is known, only additional occurrences will be scored
        - if a letter's frequency is unknown, all occurrences will be scored

        Returns:
            dict with format {letter: [occ1, occ2, occ3, occ4, occ5]}
            where the key (letter) is an upper case letter
            and the value list indicates which occurrences should be scored:
            1 = score, 0 = do not score
        """
        mask = {}

        for letter in string.ascii_uppercase:

            # Is there a rule for this letter?
            if letter in self:
                rule = self[letter]

                # Known frequency = do not score
                if rule.count_op == "eq":
                    mask[letter] = [0, 0, 0, 0, 0]

                # Score only unknown frequencies
                else:
                    mask[letter] = []
                    for freq in range(1, 6):
                        mask[letter].append(1 if freq > rule.count else 0)

            # No rule = score all frequencies
            else:
                mask[letter] = [1, 1, 1, 1, 1]

        return mask

    def _get_position_mask(self) -> dict[str, list[int]]:
        """Gets the mask for position scores

        The mask indicates which letter positions are of relevance when scoring a word.

        Rules are;
        - if a letter's position is known, no values will be scored
        - if a letter has excluded/confirmed positions, those values will not be scored
        - if a letter's position is unknown, all values will be scored

        Returns:
            dict with format {letter: [pos1, pos2, pos3, pos4, pos5]}
            where the key (letter) is an upper case letter
            and the value list indicates which positions should be scored:
            1 = score, 0 = do not score
        """
        mask = {}

        for letter in string.ascii_uppercase:

            # Is there a rule for this letter?
            if letter in self:
                rule = self[letter]

                # Letter position known: return zeros
                if rule.count_op == "eq" and len(rule.confirmed) == rule.count:
                    mask[letter] = [0, 0, 0, 0, 0]

                # Letter position partly known: mask known values
                else:
                    mask[letter] = []
                    known = rule.confirmed | rule.excluded
                    for pos in range(5):
                        mask[letter].append(0 if pos in known else 1)

            # No rule = score all positions
            else:
                mask[letter] = [1, 1, 1, 1, 1]

        return mask

    def _get_frequency_scores(self, words: list[str]) -> dict[str, list[int]]:
        """Returns dict of letter frequency scores in word list

        Args:
            words: list of five-letter words

        Returns:
            dict with format {letter: [occ1, occ2, occ3, occ4, occ5]}
            where the key (letter) is an upper case letter
            and the value list indicates how many times that letter occurs in the word list
            occ1 counts in how many words the letter occurs at least once
            occ2 counts in how many words the letter occurs at least twice, etc.
            >>> _get_frequency_scores(["ABBEY"])
            {
                'A': [1, 0, 0, 0, 0],
                'B': [1, 1, 0, 0, 0], ...,
                'E': [1, 0, 0, 0, 0], ...,
                'Z': [0, 0, 0, 0, 0],
            }
        """

        # Get mask
        mask = self._get_frequency_mask()

        # Initialize count dict
        frequency_scores = {
            letter: [0, 0, 0, 0, 0] for letter in string.ascii_uppercase
        }

        # Iterate over word list
        for word in words:

            # Count of letter occurrences
            counter = Counter(word)

            # Increment tally for each letter in the word
            for letter, count in counter.items():

                # Increment letter[0] for first occurence, letter[1] for second etc.
                for index in range(count):

                    # Check for inclusion
                    if mask[letter][index]:

                        # Increment if included
                        frequency_scores[letter][index] += 1

        return frequency_scores

    def _get_position_scores(self, words: list[str]) -> dict[str, list[int]]:
        """Returns dict of letter position scores in word list

        Args:
            list of five-letter words

        Returns:
            dict with format {letter: [pos1, pos2, pos3, pos4, pos5]}
            where the key (letter) is an upper case letter
            and the value list indicates how many times that letter occurs in the word list
            pos1 counts in how many words the letter occurs in position 1
            pos2 counts in how many words the letter occurs in position 2, etc.
            >>> _get_position_scores(["ABBEY"])
            {'A': [1, 0, 0, 0, 0], 'B': [0, 1, 1, 0, 0], ...}
        """

        # Get mask
        mask = self._get_position_mask()

        # Initialize count dict
        position_scores = {letter: [0, 0, 0, 0, 0] for letter in string.ascii_uppercase}

        # Iterate over word list
        for word in words:

            # Increment score for each letter in the word
            for index, letter in enumerate(word):
                position_scores[letter][index] += mask[letter][index]

        return position_scores


def _score_words(
    words: list[str],
    fscores: dict[str, list[int]],
    pscores: dict[str, list[int]],
) -> list[WordScore]:
    """Score words as per given frequency and positio scores

    Args:
        words: list of words to score
        fscores: frequency score dict for each letter
        pscores; position score dict for each letter

    Returns:
        list of WordScores for this list of words
    """

    word_scores = []

    # Score each word
    for word in words:

        score = 0

        # Frequency score
        counter = Counter(word)
        for letter, count in counter.items():
            for index in range(count):
                score += fscores[letter][index]

        # Position score
        for position, letter in enumerate(word):
            score += pscores[letter][position]

        word_scores.append(WordScore(word, score))

    return word_scores
