#!/usr/bin/env -S uv run

"""
From https://www.redblobgames.com/x/2635-a-vs-an/
Copyright 2026 Red Blob Games <redblobgames@gmail.com>
@license Apache-2.0 <http://www.apache.org/licenses/LICENSE-2.0.html>

To get cmudict:

mkdir third-party
cd third-party
git clone --depth 1 https://github.com/cmusphinx/cmudict
"""

from dataclasses import dataclass, field
import re
import json
import random
import a_vs_an

# These are in /usr/share/dict/words but nth is weird and urman is a name
EXCLUDE = ("nth urman").split(" ")

def read_vowel_sounds(filename):
    vowel_sounds = set()
    with open(filename, 'r') as f:
        for line in f:
            sound, type = line.strip().split("\t")
            if type == 'vowel':
                vowel_sounds.add(sound)
    return vowel_sounds


def read_cmudict(filename):
    entries = []
    with open(filename, 'r') as f:
        for line in f:
            word, pronunciation = line.strip().split(" ", 1)
            pronunciation = re.sub(r"\d", "", pronunciation)
            entries.append([word, pronunciation])

    return entries


def find_single_letter_pronunciation(cmudict_entries):
    word_pronunciations = dict()
    for word, pronunciation in cmudict_entries:
        if len(word) != 1: continue
        word_pronunciations[word] = pronunciation

    # Special case where single letter is also a word, pronounced differently than the letter
    word_pronunciations['a'] = "EY"
    word_pronunciations['g'] = "G IY"

    return word_pronunciations


def is_initialism(word, pronunciation, single_letter):
    # Example: mit EH M AY T IY
    # Example: sms EH S EH M EH S
    # Fails: mba EH M B IY EY (added to EXCLUDE manually)
    spelled_pronunciation = " ".join([single_letter[letter] for letter in word])
    return pronunciation == spelled_pronunciation


def filter_cmudict(cmudict_entries):
    single_letter = find_single_letter_pronunciation(cmudict_entries)
    usr_share_dict_words = set([line.strip() for line in open("/usr/share/dict/words", 'r')])
    usr_share_dict_words_lowercase = set([word.lower() for word in usr_share_dict_words])

    results = []
    filtered_single_letter = 0
    filtered_secondary = 0
    filtered_punctuation = 0
    filtered_initialism = 0
    filtered_exclusions = 0
    filtered_propername = 0
    filtered_nondictionary = 0

    for word, pronunciation in cmudict_entries:
        # Single letters are not interesting to me
        if len(word) == 1:
            filtered_single_letter += 1
            continue

        # Some words have multiple pronunciations, but I'm not handling that here.
        # The ones that differ on a vs an are: (herb, herbal, etc. ; him ; homage ; replace)
        if '(' in word:
            filtered_secondary += 1
            continue

        # We just won't handle words with punctuation in this project
        if not re.match(r"^\w+$", word) or not re.match(r"^[\w\s]+$", pronunciation):
            filtered_punctuation += 1
            continue

        # Proper names aren't interesting to me, and this dictionary capitalizes them
        if word not in usr_share_dict_words:
            if word in usr_share_dict_words_lowercase:
                filtered_propername += 1
            else:
                filtered_nondictionary += 1 # some of these are proper names too but I can't separate them out
            continue

        # Initialisms like "SMS" are spelled out. Although that's a
        # valid case, I think it's not interesting for this project,
        # so I'm going to filter those out.
        if is_initialism(word, pronunciation, single_letter):
            filtered_initialism += 1
            continue

        # Some words I just don't want to handle
        if word in EXCLUDE:
            filtered_exclusions += 1
            continue

        results.append((word, pronunciation))

    print("Filtered for single letter:", filtered_single_letter)
    print("Filtered for secondary:", filtered_secondary)
    print("Filtered for punctuation:", filtered_punctuation)
    print("Filtered for proper name:", filtered_propername)
    print("Filtered for nondictionary:", filtered_nondictionary)
    print("Filtered for exclusion:", filtered_exclusions)
    print("Filtered for initialism:", filtered_initialism)
    return results


def annotate_words(cmudict_entries, vowel_sounds):
    VOWEL_LETTERS = "aeiou"

    annotated_words = []
    for word, pronunciation in cmudict_entries:
        first_letter_is_vowel = word[0] in VOWEL_LETTERS
        first_sound_is_vowel = pronunciation.split(" ")[0] in vowel_sounds

        annotated_word = ('+' if first_letter_is_vowel else '#') + word + ('+' if first_sound_is_vowel else '#')
        annotated_words.append(annotated_word)

    return annotated_words


def test_words(cmudict_entries):
    pass


# I used this trie to build visualizations
@dataclass
class Trie:
    count_vowel: int = 0
    count_consonant: int = 0
    examples = set()
    children: dict[str, Trie] = field(default_factory=dict)

    @property
    def count(self) -> int:
        return self.count_vowel + self.count_consonant

    @property
    def vowel(self) -> str:
        if self.count_vowel > 0 and self.count_consonant == 0: return 'all'
        if self.count_vowel == 0 and self.count_consonant > 0: return 'none'
        if self.count_vowel > 0 and self.count_consonant > 0: return 'some'
        raise ValueError("Empty node")

    def add(self, word: str):
        if word == '': return
        first, rest = word[0], word[1:]
        if first not in self.children: self.children[first] = Trie()
        self.children[first].add(rest)

    def analyze(self, path=""):
        self.count_vowel = 0
        self.count_consonant = 0
        self.examples = set()
        for letter, child in self.children.items():
            if letter == '+':
                child.count_vowel = 1
                child.examples = set([path])
            elif letter == '#':
                child.count_consonant = 1
                child.examples = set([path])
            else:
                child.analyze(path + letter)

            self.count_vowel += child.count_vowel
            self.count_consonant += child.count_consonant

        if self.count_vowel == 0 or self.count_consonant == 0:
            for child in self.children.values():
                self.examples |= child.examples

    def to_json(self, letter='', depth=1):
        array = []
        if depth > 0:
            for child_letter, child in self.children.items():
                array.append(child.to_json(child_letter, depth-1))
        json = {'name': letter, 'vowel': self.vowel, 'children': array}
        return json

    def to_truncated_json(self, letter='', prefix=''):
        # TODO: this seems like it should be more compact; investigate

        def format_set(s):
            elements = list(s)
            random.shuffle(elements)
            return ", ".join(elements[:4])

        array = []
        examples = self.examples
        if self.count == 1:
            # collapse all the nodes
            name = prefix
            children = self.children.items()
            while children:
                first_child = list(children)[0]
                if len(children) > 1: raise ValueError("Multiple children")
                name += first_child[0]
                examples = first_child[1].examples
                children = first_child[1].children.items()
        elif self.vowel == 'all' or self.vowel == 'none':
            name = prefix + "*" # leaf
        else:
            name = letter # internal node
            default_type = ''
            if self.count_vowel > self.count_consonant: default_type = 'all'
            if self.count_consonant >= self.count_vowel: default_type = 'none'

            for child_letter, child in self.children.items():
                if child.vowel == default_type:
                    examples |= child.examples
                else:
                    array.append(child.to_truncated_json(child_letter, prefix+child_letter))

            grouped_node = {'name': prefix + "*", 'example': format_set(examples), 'vowel': default_type, 'count_vowel': self.count_vowel, 'count_consonant': self.count_consonant, 'children': []}
            if default_type != '': array.append(grouped_node)

        name = name.replace("+", "").replace("#", "")
        json = {'name': name, 'example': format_set(examples), 'count': self.count, 'vowel': self.vowel, 'count_vowel': self.count_vowel, 'count_consonant': self.count_consonant, 'children': array}
        return json


def main():
    vowel_sounds = read_vowel_sounds("third-party/cmudict/cmudict.phones")
    cmudict_entries = read_cmudict("third-party/cmudict/cmudict.dict")
    filtered_entries = filter_cmudict(cmudict_entries)
    annotated_words = annotate_words(filtered_entries, vowel_sounds)
    print("Words to be analyzed:", len(filtered_entries))

    trie = Trie()
    for word in annotated_words: # first character is letter-is-vowel, last is sound-is-vowel
        trie.add(word[1:])

    # Print all mismatches between the incorrect rule (first letter) and correct rule (first sound)
    with open("_build/approximation-0.txt", 'w') as f:
        for word in annotated_words:
            if word[0] != word[-1]:
                f.write(f"{word[0]} {word[1:-1]} {word[-1]} {[p for [w, p] in filtered_entries if w == word[1:-1]]}\n")

    # Print all mismatches between the approximate rule (from a_vs_an.py) and correct rule (first sound)
    with open("_build/approximation-1.txt", 'w') as f:
        for word in annotated_words:
            is_vowel_guess = a_vs_an.starts_with_vowel(word[1:-1])
            is_vowel_actual = word[-1] == '+'
            if is_vowel_guess != is_vowel_actual:
                f.write(f"{word[0]} {word[1:-1]} {is_vowel_actual} {is_vowel_guess} {[p for [w, p] in filtered_entries if w == word[1:-1]]}\n")

    trie.analyze()

    # This is a nice visualization showing whether two letters is
    # enough. It is a lot of the time but for some letter combinations
    # we need more.
    json.dump(trie.to_json('*', 2), open("_build/depth-2.json", 'w'), indent=2)

    # This goes all the way to the last character needed to make the decision
    json.dump(trie.to_truncated_json('*'), open("_build/truncated.json", 'w'), indent=2)


if __name__ == '__main__':
    main()