from LanguageBase import Language 
from taibun import Converter 
from transformers import pipeline 
model_repo_id = "emlinking/wav2vec2-large-xls-r-300m-tsm-asr-v6" 
util = Converter() 
def get_initial_length(s):
    if s[0] in ['a', 'e', 'i', 'o', 'u']: return 0 
    if (len(s) == 3 and s[:2] == 'ng') or (len(s) == 2 and s[:1] == 'm'): return 0  
    if len(s) >= 3 and s[2] == 'h':
        return 3
    elif len(s) >= 2 and s[1] in ['s', 'h']:
        return 2
    else:
        return 1 
class TaiwaneseHokkien(Language):
    def __init__(self, **kwargs):
        self.pipe = pipeline(task="automatic-speech-recognition", model=model_repo_id, **kwargs)
    def asr(self, audio):
        return self.pipe(audio)['text']
    def compare(self, target_pron, user_pron):
        # normalize 
        target_pron = target_pron.lower().split()
        user_pron = user_pron.lower().split() 
        result = [] 
        for i in range(min(len(target_pron), len(user_pron))):
            target_syls = [x for x in target_pron[i].split('-') if x]
            user_syls = [x for x in user_pron[i].split('-') if x]
            for j in range(min(len(target_syls), len(user_syls))):
                target_syl = util._Converter__get_number_tone(target_syls[j])
                user_syl = util._Converter__get_number_tone(user_syls[j])
                til = get_initial_length(target_syl)
                uil = get_initial_length(user_syl)
                if target_syl[:til] != user_syl[:uil]:
                    if uil == 0:
                        result.append((' ', 'initial error'))
                    else:
                        result.append((user_syl[:uil], 'initial error'))
                else:
                    result.append((user_syl[:uil], None))
                if target_syl[til:-1] != user_syl[uil:-1]:
                    result.append((user_syl[uil:-1], 'rime error'))
                else:
                    result.append((user_syl[uil:-1], None))
                if target_syl[-1] != user_syl[-1]:
                    result.append((user_syl[-1], 'tone error'))
                else:
                    result.append((user_syl[-1], None))
                if j < min(len(target_syls), len(user_syls))-1:
                    result.append(('-', None))
            if len(target_syls) > len(user_syls):
                for syl in target_syls[len(user_syls):]:
                    result.append(('-' + syl, 'missing syllables'))
            elif len(user_syls) > len(target_syls):
                for syl in user_syls[len(target_syls):]:
                    result.append(('-' + syl, 'extra syllables'))
            result.append((' ', None))
        if len(target_pron) > len(user_pron):
            for word in target_pron[len(user_pron):]:
                result.append((word, 'missing syllables'))
                result.append((' ', None))
        elif len(user_pron) > len(target_pron):
            for word in user_pron[len(target_pron):]:
                result.append((word, 'extra syllables'))
                result.append((' ', None))
        return result 
    @property 
    def compare_colors(self):
        return {'tone error': 'red', 'initial error': 'blue', 'rime error': 'green', 'missing syllables': 'yellow', 'extra syllables': 'lime'}