import os
import json

_special_tokens_phonemes = list('!"(),./:;?[]{}¡ª°´º¿\'')
_special_tokens_graphemes = ['a', 'á', 'e', 'é', 'i', 'í', 'o', 'ó', 'u', 'ú', 'b', 'c', 'ch', 'd', 'f', 'g', 'gu',
                             'h', 'j', 'k', 'l', 'll', 'm', 'n', 'nh', 'ñ', 'p', 'q', 'qu', 'r', 'rr', 's', 't', 'v', 'x', 'z']

# Export all symbols:

_module_dir = os.path.dirname(__file__)
_json_path = os.path.join(_module_dir, 'phoneme_token_maps.json')

symbols_phonemes = []
symbols_graphemes = []

if os.path.exists(_json_path):
    try:
        with open(_json_path, 'r', encoding='utf-8') as fh:
            j = json.load(fh)

        for k, v in j.items():
            phoneme = v.get('phoneme')
            grapheme = v.get('grapheme')
            if phoneme not in symbols_phonemes:
                symbols_phonemes.append(phoneme)
            if grapheme not in symbols_graphemes:
                symbols_graphemes.append(grapheme)
        for special_token in _special_tokens_phonemes:
            if special_token not in symbols_phonemes:
                symbols_phonemes.append(special_token)
        for special_token in _special_tokens_graphemes:
            if special_token not in symbols_graphemes:
                symbols_graphemes.append(special_token)

    except Exception as e:
        print(f"Error loading phoneme_token_maps.json: {e}")

# Create mapping dictionaries
symbol_to_id_phonemes = {symbol: i for i,
                         symbol in enumerate(symbols_phonemes)}
id_to_symbol_phonemes = {i: symbol for i,
                         symbol in enumerate(symbols_phonemes)}

symbol_to_id_graphemes = {symbol: i for i,
                          symbol in enumerate(symbols_graphemes)}
id_to_symbol_graphemes = {i: symbol for i,
                          symbol in enumerate(symbols_graphemes)}

print("Phoneme symbols:\n")
for i, s in enumerate(symbols_phonemes):
    print(f"{i:3d}: {s}")
print("Grapheme symbols:\n")
for i, s in enumerate(symbols_graphemes):
    print(f"{i:3d}: {s}")


TOKEN_SIZE = len(symbols_phonemes)
VOCAB_SIZE = len(symbols_graphemes)


class TextCleanerGal:
    def __init__(self):
        self.phoneme_to_id = symbol_to_id_phonemes
        self.grapheme_to_id = symbol_to_id_graphemes

        self.id_to_phoneme = id_to_symbol_phonemes
        self.id_to_grapheme = id_to_symbol_graphemes

    def __call__(self, tokens, mode="phoneme"):
        if mode == "phoneme":
            dictionary = self.phoneme_to_id
            unk = "X"
        else:
            dictionary = self.grapheme_to_id
            unk = " "
        return [dictionary.get(tok, dictionary[unk]) for tok in tokens]

    def decode_phoneme(self, idxs):
        return "".join(self.id_to_phoneme.get(i, "X") for i in idxs)

    def decode_grapheme(self, idxs):
        return "".join(self.id_to_grapheme.get(i, "X") for i in idxs)