from collections import abc
import re

from pudzu.utils import *

# Pronouncing dictionary class, largely targeted at English

IPA_VOWELS = "ɔɑiuɛɪʊʌəæeaoyʏøɝɚɒɜʁ"
IPA_VOWELENDINGS = IPA_VOWELS + "ː̯̩̃"
IPA_CONSONANTS = "pbtdkɡfvθðszʃʒhmnŋlrɾɹjwʔʍχ"
IPA_STRESS = "ˈˌ"

class Nouncer(abc.MutableMapping):

    PHONEME_PATTERN = ("("
        "[{stress}]|"
        "[{vowel}][ː̃]*[ʊɪə]?[̯]?|"
        "(?:t͡?ʃ|d͡?ʒ|[{consonant}])[̩]?"
    ")[.ˑ ]?".format(stress=IPA_STRESS, vowel=IPA_VOWELS, consonant=IPA_CONSONANTS))

    ARPABET_TO_IPA = {
        'AO': 'ɔ', 'AA': 'ɑ', 'IY': 'i', 'UW': 'u', 'EH': 'ɛ', 'IH': 'ɪ', 'UH': 'ʊ', 'AH': 'ʌ', 'AX': 'ə', 'AE': 'æ',
        'EY': 'eɪ', 'AY': 'aɪ', 'OW': 'oʊ', 'AW': 'aʊ', 'OY': 'ɔɪ', 'ER': 'ɝ', 'AXR': 'ɚ',
        'P': 'p', 'B': 'b', 'T': 't', 'D': 'd', 'K': 'k', 'G': 'ɡ', 'CH': 't͡ʃ', 'JH': 'd͡ʒ', 'F': 'f', 'V': 'v',
        'TH': 'θ', 'DH': 'ð', 'S': 's', 'Z': 'z', 'SH': 'ʃ', 'ZH': 'ʒ', 'HH': 'h', 'M': 'm', 'EM': 'm̩',
        'N': 'n', 'EN': 'n̩', 'NG': 'ŋ', 'ENG': 'ŋ̍', 'L': 'l', 'EL': 'ɫ̩', 'R': 'r', 'DX': 'ɾ',
        'Y': 'j', 'W': 'w', 'Q': 'ʔ' }
    
    ENPR_TO_IPA = { 
        "GenAm": { 'ä': 'ɑ', 'ă': 'æ', 'ǎ': 'æ', 'ā': 'eɪ', 'ĕ': 'ɛ', 'ē': 'i', 'i': 'ɪ', 'ĭ': 'ɪ', 'ī': 'aɪ',
        'ŏ': 'ɑ', 'ō': 'oʊ', 'ô': 'ɔ', 'ŭ': 'ʌ', 'ə': 'ə',
        'oi': 'ɔɪ', 'ŏŏ': 'ʊ', 'o͝o': 'ʊ', 'ōō': 'u', 'o͞o': 'u', 'ou': 'aʊ',
        'ăr': 'æɹ', 'är': 'ɑɹ', 'âr': 'ɛɹ', 'ĕr': 'ɛɹ', 'ĭr': 'ɪɹ', 'îr': 'iɹ',
        'ōr': 'ɔɹ', 'ôr': 'ɔɹ', 'ŏŏr': 'ʊɹ', 'o͝or': 'ʊɹ', 'ûr': 'ɝ', 'ər': 'ɚ',
        'b': 'b', 'ch': 't͡ʃ', 'd': 'd', 'f': 'f', 'g': 'ɡ', 'h': 'h', 'hw': 'ʍ', 'j': 'd͡ʒ', 'k': 'k',
        'l': 'l', 'm': 'm', 'n': 'n', 'ng': 'ŋ', 'p': 'p', 'r': 'r', 's': 's', 'sh': 'ʃ', 't': 't', 'th': 'θ',
        'v': 'v', 'w': 'w', 'y': 'j', 'z': 'z', 'zh': 'ʒ' } }
    
    def __init__(self, filename=None, normalize=str.lower):
        self.pdict = KeyEquivalenceDict(normalizer=normalize)
        if filename is not None:
            with open(filename, "r", encoding="utf-8") as f:
                for entry in f:
                    m = re.match("(.*)\t(.*)", entry)
                    if m:
                        self.pdict.setdefault(m.group(1), set()).add(tuple(m.group(2).split(" ")))
        
    def save(self, filename):
        """Save pronuncing dictionary to file."""
        with open(filename, 'w', encoding='utf-8') as f:
            for word in sorted(self.pdict.keys()):
                for pronunciation in self.pdict[word]:
                    print("{}\t{}".format(word, " ".join(pronunciation)), file=f)
     
    def import_cmudict(self, filename):
        """Import pronunciations from CMUDict."""
        with open(filename, "r", encoding="latin-1") as f:
            for entry in f:
                m = re.match("([^(]*)(?:[(][0-9][)])?  (.*)", entry)
                if m:
                    self.pdict.setdefault(m.group(1).lower(), set()).add(self.arpabet_to_phonemes(m.group(2)))
        
    def import_list(self, filename, delimiter="\t", error_cache=None, enpr=False):
        """Import pronunciations from a file, either in IPA or enPR format."""
        if error_cache: error_cache.set([])
        to_phonemes = partial(self.enpr_to_phonemes, dialect=enpr) if enpr else self.ipa_to_phonemes
        with open(filename, "r", encoding="utf-8") as f:
            for entry in f:
                m = re.match("(.*)[{}](.*)".format(delimiter), entry)
                if m:
                    try:
                        self.pdict.setdefault(m.group(1), set()).update(to_phonemes(m.group(2)))
                    except ValueError as e:
                        if not error_cache: 
                            raise ValueError("Error reading {}: {}".format(m.group(1), str(e)))
                        error_cache.value.append((m.group(1), str(e)))
        
    PHONEME_REGEX = re.compile(PHONEME_PATTERN)
    PRONUNCIATION_REGEX = re.compile("(?:{phoneme})*".format(phoneme=PHONEME_PATTERN))
    PARENTHESES_REGEX = re.compile("[(]([^)]+)[)]")
    BRACKETS_REGEX = re.compile("\[([^\]]+)\]")

    def ipa_to_phonemes(self, pronunciation):
        def normalise(string):
            string = re.sub("(?<=.)-(?=.)", "", string)
            return (string.replace('ĭ','i').replace('̈','').replace('ᵻ','[ɪ,ə]')
                          .replace('ɨ','ɪ').replace('ɘ','ə').replace('ɵ','ʊ')
                          .replace('ɱ','m'))
        def expand_parentheses(string):
            m = re.search(self.PARENTHESES_REGEX, string)
            if not m: return expand_brackets(string)
            else: return [x for s in (re.sub(self.PARENTHESES_REGEX, "", string, count=1), re.sub(self.PARENTHESES_REGEX, r"\1", string, count=1)) for x in expand_parentheses(s)]
        def expand_brackets(string):
            m = re.search(self.BRACKETS_REGEX, string)
            if not m: return [string]
            else: return [x for s in (re.sub(self.BRACKETS_REGEX, v, string, count=1) for v in m.group(1).split(",")) for x in expand_brackets(s)]
        pronunciations = expand_parentheses(normalise(pronunciation))
        return [self.ipa_to_phonemes_no_parentheses(p) for p in pronunciations]
        
    def ipa_to_phonemes_no_parentheses(self, pronunciation):
        m = re.match(self.PRONUNCIATION_REGEX, pronunciation)
        if m.end() < len(pronunciation):
            raise ValueError("Invalid pronunciation: /{}**{}/".format(pronunciation[:m.end()], pronunciation[m.end():]))
        phonemes, stress, vowels = [], "", 0
        for p in re.findall(self.PHONEME_REGEX, pronunciation):
            if p in IPA_STRESS:
                stress = p
            elif p[0] in IPA_VOWELS:
                phonemes.append(stress+p)
                stress = ""
                vowels += 1
            else:
                if p == "tʃ": p = "t͡?ʃ"
                elif p == "dʒ": p = "d͡?ʒ"
                phonemes.append(p)
        if vowels == 1: 
            try: # autoadd stress mark for monosyllabic words
                i = next(i for i in range(len(phonemes)) if phonemes[i][0] in IPA_VOWELS)
                if phonemes[i] != 'ə':
                    phonemes[i] = 'ˈ' + phonemes[i]
            except StopIteration:
                pass
        return tuple(phonemes)

    def arpabet_to_phonemes(self, pronunciation):
        phonemes = []
        for l in pronunciation.split(" "):
            phoneme = ""
            if l == 'AH0': l = 'AX0' # distinguish ʌ from ə
            if l[-1] in ('0', '1', '2'):
                phoneme += { '0': '', '1': 'ˈ', '2': 'ˌ' }[l[-1]]
                l = l[:-1]
            phoneme += self.ARPABET_TO_IPA[l]
            phonemes.append(phoneme)
        return tuple(phonemes)
        
    def enpr_to_phonemes(self, pronunciation, dialect):
        i, vi, phonemes, tdict = 0, 0, [], self.ENPR_TO_IPA[dialect]
        while i < len(pronunciation):
            for l in (3, 2, 1):
                if pronunciation[i:i+l] in tdict:
                    phonemes.append(tdict[pronunciation[i:i+l]])
                    if self._is_vowel(phonemes[-1]): vi = len(phonemes) - 1
                    i += len(pronunciation[i:i+l])
                    break
            else:
                if pronunciation[i] in " -'ʹ′":
                    if pronunciation[i] in "'ʹ′": phonemes.insert(vi, 'ˈ')
                    i += 1
                elif pronunciation[i] in "()":
                    phonemes.append(pronunciation[i])
                    i += 1
                else:
                    raise ValueError("Unrecognised enPR character {} in {}".format(pronunciation[i], pronunciation))
        return self.ipa_to_phonemes("".join(phonemes))
        
    @classmethod
    def _is_vowel(self, phoneme): return phoneme[-1] in IPA_VOWELENDINGS
    
    @classmethod
    def _is_stressed(self, phoneme, secondary=False): return phoneme[0] == 'ˈ' or secondary and phoneme[0] == 'ˌ'
    
    @classmethod
    def _is_consonant(self, phoneme): return phoneme[-1] not in IPA_VOWELENDINGS
    
    # MutableMapping
    
    def __getitem__(self, word):
        """Return the pronunciations of a word in IPA. Syllable boundaries aren't marked and stress marks go before the stressed vowel."""
        return {"".join(pronunciation) for pronunciation in self.pdict[word]}
        
    def __setitem__(self, word, pronunciations):
        """Set pronunciations of a word in IPA. If a single pronunciation is given then it is appended to existing ones."""
        if non_string_iterable(pronunciations):
            self.pdict[word] = { q for p in pronunciations for q in self.ipa_to_phonemes(p) }
        else:
            self.pdict.setdefault(word, set()).update(self.ipa_to_phonemes(pronunciations))
        
    def __delitem__(self, word):
        """Delete all pronunciations of a word."""
        del self.pdict[word]
        
    def __iter__(self):
        return iter(self.pdict)
        
    def __len__(self):
        return len(self.pdict)
        
    def __repr__(self):
        return "<Nouncer: {} entries>".format(len(self))
        
    # rest of API
    
    def pronunciations(self, word_filter=None, pronunciation_filter=None):
        """Generator returning matching pronunciations."""
        def match(pattern, string):
            return True if pattern is None else pattern(string) if callable(pattern) else re.search(pattern, string)
        return ((w, p) for w,ps in self.items() if match(word_filter,w) for p in ps if match(pronunciation_filter, p))
    
    def syllables(self, word, default_counter=None):
        """Number of syllables in a word."""
        try:
            return {"".join(pronunciation) : sum(1 for phoneme in pronunciation if self._is_vowel(phoneme)) for pronunciation in self.pdict[word]}
        except KeyError:
            if default_counter is not None:
                return { "({})".format(word): default_counter(word) }
            raise
       
    @classmethod
    def _rhymeswith(self, phonemes1, phonemes2, identirhyme=False, cutrhyme=False, multirhyme=False):
        stress1 = first((i for i in range(len(phonemes1)) if self._is_stressed(phonemes1[i])), 0)
        stress2 = first((i for i in range(len(phonemes2)) if self._is_stressed(phonemes2[i])), 0)
        same_consonant = stress1==stress2==0 or stress1>0 and stress2>0 and phonemes1[stress1-1]==phonemes2[stress2-1]
        pattern1 = phonemes1[stress1:len(phonemes2)-stress2+stress1 if cutrhyme else None]
        pattern2 = phonemes2[stress2:]
        if multirhyme:
            def strip_consonants(p):
                lv = first((i for i in reversed(range(len(p))) if self._is_vowel(p[i])), 0)
                return [x for i,x in enumerate(p) if self._is_vowel(x) or i >= lv]
            pattern1 = strip_consonants(pattern1)
            pattern2 = strip_consonants(pattern2)
        return pattern1 == pattern2 and (not same_consonant or identirhyme)
        
    def rhymes(self, word, identirhyme=False, cutrhyme=False, multirhyme=False):
        """Rhymes for a word."""
        d = {}
        for p1 in self.pdict[word]:
            for w2, ps2 in self.pdict.items():
                for p2 in ps2:
                    if self._rhymeswith(p1, p2, identirhyme=identirhyme, cutrhyme=cutrhyme, multirhyme=multirhyme):
                        d.setdefault("".join(p1), []).append(w2)
        return d
        
def english_syllables(word):
    """Simple heuristic for the number of syllables in an English word. 91% agreement with CMUDict."""
    pos = ["[aeiouy]+",
           "[^cgj]eo|[^cgst]ia|ii|[^cgstx]io|io$|[^g]iu|[^qg]ua|[^g]uo",
           "^mc|(s|th)ms?$",
           "[aeiouy]ing"]
    neg = ["[aeiouy]n?[^aeiouy]h?e$",
           "[aeiouy]([^aeiouytd]|tc?h)+ed$",
           "[aeiouy]r?[^aeiouycszxh]h?es$",
           "cally$|[^ei]ely$"]
    return sum(len(re.findall(r, word)) for r in pos) - sum(len(re.findall(r, word)) for r in neg)

def extract_from_wiktionary(input, output, language="en", accents=("US", "USA", "GA", "GenAm", None), enpr=False):
    """Extract pronunciations from wiktionary xml dump."""
    title_regex = re.compile("<title>(.*)</title>")
    ipa_regex_1 = re.compile("{{{{IPA[|]/([^|]+)/[|]lang={lang}}}}}".format(lang=language))
    ipa_regex_2 = re.compile("{{{{IPA[|]lang={lang}[|]/([^|]+)/}}}}".format(lang=language))
    enpr_regex = re.compile("{{enPR[|]([^}]+)}}")
    accent_regex = re.compile("{{{{a(ccent)?[|]([^}}]+[|])?({accents})[|}}]".format(accents="|".join(a for a in make_iterable(accents) if a)))
    any_accent_regex = re.compile("{{a(ccent)?[|]")
    match = ValueBox()
    with open(input, "r", encoding="utf-8") as i:
        with open(output, "w", encoding="utf-8") as o:
            for line in i:
                if match << re.search(title_regex, line):
                    title = match().group(1)
                elif ((enpr and match << re.search(enpr_regex, line)) or 
                      (not enpr and match << (re.search(ipa_regex_1, line) or re.search(ipa_regex_2, line)))):
                    if accents and not re.search(accent_regex, line) and (None not in accents or re.search(any_accent_regex, line)):
                        continue
                    elif ":" in title:
                        continue
                    for pronunciation in match().group(1).split("|" if enpr else ", "):
                        print("{}\t{}".format(title, pronunciation), file=o)
