import sys,os,re
import phonecodes
import phonecode_tables

LEAF='LEAF'
MAXSET=8000

def dict_enter_word(d,w,p):
    if len(p) == 0:
        if LEAF not in d:
            d[LEAF] = set()
        d[LEAF].add(w)
    else:
        if p[0] not in d:
            d[p[0]] = {}
        dict_enter_word(d[p[0]],w,p[1:])

def dict_convert_phonecode(d0,code0,code1,language):
    d1 = dict()
    for p0 in d0:
        if p0==LEAF:
            d1[p0]=d0[p0]
        else:
            p1 = phonecodes.convert(p0,code0,code1,language)
            d1[p1] = dict_convert_phonecode(d0[p0],code0,code1,language)
    return(d1)

def dict_phones2words(d, rtd, ph, dist=0, paths=[[]], nmax=MAXSET):
    '''Return up to nmax wordstrings, in a list, that match ph
    with up to `dist' distance.'''
    res = []
    # If we're at the root node, with no ph, then current paths is a result
    # Prepend each path with the remaining dist
    if d == rtd and len(ph)==0:
        res.extend([[-dist]+x for x in paths])
    # If the next phone can extend current word, then do it
    if len(ph)>0 and ph[0] in d:
        res.extend(dict_phones2words(d[ph[0]],rtd,ph[1:],dist,paths,nmax))
    # If we've finished any word, then try to start a new one
    if d != rtd and LEAF in d:
        for w in d[LEAF]:
            ws = [ x+[w] for x in paths ]
            r=dict_phones2words(rtd,rtd,ph,dist,ws,nmax)
            while(len(res)<nmax and len(r)>0):
                res.append(r.pop())
    # If we can sub the next phone to extend current word, then do it
    if len(ph)>0 and ph[0] not in d and dist > 0:
        for p in d:
            if p != LEAF:
                res.extend(dict_phones2words(d[p],rtd,ph[1:],dist-1,paths,nmax))
    # If we can extend current word by skipping a dict phone, then do it
    if dist > 0:
        for p in d:
            if p != LEAF:
                res.extend(dict_phones2words(d[p],rtd,ph,dist-1,paths,nmax))
    # If we can extend current word by skipping a listed phone, then do it
    if len(ph)>1 and ph[1] in d and dist > 0:
        res.extend(dict_phones2words(d[ph[1]],rtd,ph[2:],dist-1,paths,nmax))
    # return up to nmax results
    res = sorted(res)
    return(res[:nmax])

###########################################################
def read_raw_dictfile(filename,language):
    '''Read from unformatted dictfile: eliminate extra spaces and parens, 
    insert a tab.  Return a list of word,pron pairs.'''
    S = []
    with open(filename) as f:
        for line in f:
            words = re.split(r'\s+',line.rstrip())
            S.append((words[0], words[1:]))
    return(S)

def read_isle_dictfile(filename,language,params):
    '''Read from ISLE-formatted dictfile, which is just IPA pronunciations,
    but word has parentheses and part of speech to be removed.
    If params['discard_phones'] is a set of chars, then any phones in the set
    will be discarded.  For example, set('#.') discards morph, syl boundaries.
    If params['discard_diacritics'] is a strings, then chars in the string
    will be discarded.  For example, 'ˈˌ' discards stress markers.
    '''
    S = []
    with open(filename) as f:
        for line in f:
            txt=re.sub(r'\([^\)]*\)','',line.rstrip())
            if 'discard_diacritics' in params:
                patstr=r'[%s]+'%(params['discard_diacritics'])
                txt = re.sub(patstr,'',txt)
            words = re.split(r'\s+',txt)
            if 'discard_phones' in params:
                pr=[p for p in words[1:] if p not in params['discard_phones']]
            else:
                pr = words[1:]
            S.append((words[0], pr))
    return(S)

_babel_pcols = {
    'amh':2,'asm':2,'ben':2,'yue':2,'ceb':1,'luo':1,'kat':2,'gug':1,'hat':1,
    'ibo':1,'jav':1,'kur':1,'lao':2,'lit':1,'mon':2,'pus':2,'swa':1,'tgl':1,
    'tam':2,'tpi':1,'tur':1,'vie':1,'zul':1
}

def read_babel_dictfile(filename, lang, params):
    '''Read from a Babel dictfile in language lang.
    If params['pcol'], then it specifies which column contains phones.
    Otherwise, try to guess based on language.
    '''
    S = []
    if 'pcol' in params:
        pcol = params['pcol']
    else:
        pcol = _babel_pcols[lang]
    with open(filename) as f:
        for line in f:
            words = re.split(r'\t',line.rstrip(), pcol)
            options = words[pcol].split('\s+')
            for option in options:
                S.append((words[0], option.split()))
    return(S)

###########################################################
_celex_pcols = { 'eng':6, 'deu':4, 'nld':4 }

def read_celex_dictfile(filename, lang, params):
    '''Read from a CELEX dictfile in language lang.'
    If params['pcol'] exists, it should be an integer,
    specifying which column contains the phonemes.
    '''
    S = []
    if 'pcol' in params:
        pcol = params['pcol']
    else:
        pcol = _celex_pcols[lang]
    with open(filename) as f:
        for line in f:
            words = re.sub(r'-','',line).rstrip().split('\\')
            word = re.sub(r'\s+','_',words[1])
            options = words[pcol].split('\\')
            for option in options:
                ol=phonecodes.attach_tones_to_vowels(list(option),set("'"),phonecode_tables._disc_vowels,1,-1)
                S.append((word, ol))
    return(S)

###################################################################
#_callhome_columns={'arz':(1,2,3,r'([@aiu%AIOUE])'),'cmn':(0,3,2,r'([iI%eEU&a@o>uR])'),'spa':(0,2,3,r'([aeiou]+)')}
_callhome_columns={'arz':(1,2,3),'cmn':(0,3,2),'spa':(0,2,3)
}
def read_callhome_dictfile(filename, L, params):
    '''Read a callhome dictionary in language L.
    If params['callhome_columns'] exists, then it should be a tuple
    (grapheme column, phone column, tone column).
    '''
    S = []
    if 'callhome_columns' in params:
        (gcol, pcol, tcol) = params['callhome_columns']
    else:
        (gcol, pcol, tcol) = _callhome_columns[L]
    with open(filename) as f:
        for line in f:
            fields = line.rstrip().split('\t')  # Split fields by tab
            tone_options = fields[tcol].split('//') # split options by slash
            syl_options = fields[pcol].split('//') 
            if len(tone_options)==1:
                options=[ (tone_options[0],s) for s in syl_options ]
            elif len(tone_options)==len(syl_options):
                options=[ (tone_options[n],syl_options[n])
                          for n in range(0,len(syl_options)) ]
            else:
                raise KeyError('tone_options and syl_options mismatched: {}'.
                               format(line.rstrip()))
            for option in options:
                syls = option[1].split(' ')
                tones = list(option[0])
                phones=[ syls[n]+tones[n] for n in range(0,len(tones)) ]
                phones.append(''.join(syls[len(tones):]))
                il = list(''.join(phones))
                ol = phonecodes.attach_tones_to_vowels(il, set(phonecode_tables._tone2ipa[L].keys()), phonecode_tables._callhome_vowels[L], -1, 1)
                S.append((fields[gcol], ol))
    return(S)

###########################################################################
def write_dictfile(w2p, filename):
    '''Write a w2p dictionary to filename'''
    lines2write = set()  # make sure each written line is unique
    for (w,pron) in w2p.items():
        if len(w)>0 and len(pron)>0:
            pl = [ re.sub(r"\s*Í¡\s*",'',p)
                   for p in pron if not p.isspace() and len(p)>0 ]
            pstr = ' '.join(pl)
            if len(pstr) > 0:
                lines2write.add('%s\t%s\n' % (w,pstr))
    with open(filename,'w') as f:
        for line in sorted(lines2write):
            f.write(line)

###################################################################
_dictread = {
    'babel':read_babel_dictfile,
    'callhome':read_callhome_dictfile,
    'celex':read_celex_dictfile,
    'isle':read_isle_dictfile
}

class lex:
    def __init__(self,language,phonecode='ipa'):
        self.p2w = {}
        self.w2p = {}
        self.language = language
        self.phonecode = phonecode
    def copy(self):
        other = lex(self.language,self.phonecode)
        other.p2w.update(self.p2w)
        other.w2p.update(self.w2p)
        return(other)
    def add(self,word,pron):
        dict_enter_word(self.p2w,word,pron)
        self.w2p[word] = pron
        return(self)
    def recode(self, code1):
        if self.phonecode==code1:
            return(self.copy())
        else:
            other = lex(self.language, code1)
            other.p2w=dict_convert_phonecode(self.p2w,self.phonecode,code1,self.language)
            for (w,p) in self.w2p.items():
                q=phonecodes.convertlist(p,self.phonecode,code1,self.language)
                other.w2p[w] = q
            return(other)
    def phones2words(self, ph, D=0, paths=[[]], nmax=MAXSET):
        '''Find all word sequences matching ph with up to D string edit dist.
        Return: a dict mapping distances to lists of lists:
        L=lex()
        res=L.phones2words(ph, 2)
        sntce in res[0] is a list of words matching ph with 0 edit distance
        sntce in res[1] is a list of words matching ph with 1 edit distance
        sntce in res[2] is a list of words matching ph with 2 edit distance
        If the total number of candidates is greater than nmax,
        then the set is truncated at nmax candidates, with a priority
        queue behavior, i.e., the highest-edit-distance candidates
        are discarded first.'''
        retval={}
        listlist=dict_phones2words(self.p2w,self.p2w,ph,D,paths,nmax)
        for res in listlist:
            if len(res)>1:
                if res[0]+D not in retval:
                    retval[res[0]+D]=[]
                retval[res[0]+D].append(res[1:])
        return(retval)
    def words2phones(self, words):
        pron = []
        for w in words:
            if w in self.w2p:
                pron.extend(self.w2p[w])
            elif w.upper() in self.w2p:
                pron.extend(self.w2p[w.upper()])
            elif w.lower() in self.w2p:
                pron.extend(self.w2p[w.lower()])
            elif w.capitalize() in self.w2p:
                pron.extend(self.w2p[w.capitalize()])
            else:
                pron.append(w)
        return(pron)
    
    def read(self,filename,lang,dicttype, params={}):
        if dicttype in _dictread:
            plist = _dictread[dicttype](filename,lang,params)
            for p in plist:
                self.add(p[0],p[1])
        else:
            plist = read_raw_dictfile(filename,lang)
            for p in plist:
                self.add(p[0],p[1])
        return(self)
    def save(self,filename):
        write_dictfile(self.w2p, filename)

_dict2phonecode = {'babel':'xsampa','celex':'disc','isle':'ipa'}
def read(filename,lang,dicttype,params={}):
    if dicttype in _dict2phonecode:
        newlex = lex(lang,_dict2phonecode[dicttype])
    else:
        newlex = lex(lang,dicttype)
    return(newlex.read(filename,lang,dicttype,params))
        
