"""adapted from https://github.com/keithito/tacotron"""

import re

_alt_re = re.compile(r"\([0-9]+\)")


class Grapheme2PhonemeDictionary:
    """Thin wrapper around g2p data."""

    def __init__(self, file_or_path, keep_ambiguous=True, encoding="latin-1"):
        with open(file_or_path, encoding=encoding) as f:
            entries = _parse_g2p(f)
        if not keep_ambiguous:
            entries = {word: pron for word, pron in entries.items() if len(pron) == 1}
        self._entries = entries

    def __len__(self):
        return len(self._entries)

    def lookup(self, word):
        """Returns list of pronunciations of the given word."""
        return self._entries.get(word.upper())


def _parse_g2p(file):
    g2p = {}
    for line in file:
        if len(line) and (line[0] >= "A" and line[0] <= "Z" or line[0] == "'"):
            parts = line.split("  ")
            word = re.sub(_alt_re, "", parts[0])
            pronunciation = parts[1].strip()
            if word in g2p:
                g2p[word].append(pronunciation)
            else:
                g2p[word] = [pronunciation]
    return g2p