from jamo import hangul_to_jamo
import librosa
import torch

sample_rate = 22050
preemphasis = 0.97
n_fft = 1024
hop_length = 256
win_length = 1024
ref_db = 20
max_db = 100
mel_dim = 80

PAD = '_'
EOS = '~'
SPACE = ' '

JAMO_LEADS = "".join([chr(_) for _ in range(0x1100, 0x1113)])
JAMO_VOWELS = "".join([chr(_) for _ in range(0x1161, 0x1176)])
JAMO_TAILS = "".join([chr(_) for _ in range(0x11A8, 0x11C3)])
ETC = ".!?"

VALID_CHARS = JAMO_LEADS + JAMO_VOWELS + JAMO_TAILS + SPACE + ETC
symbols = PAD + EOS + VALID_CHARS

_symbol_to_id = {s: i for i, s in enumerate(symbols)}
_id_to_symbol = {i: s for i, s in enumerate(symbols)}

# text를 초성, 중성, 종성으로 분리하여 id로 반환하는 함수
def text_to_sequence(text):
    sequence = []
    if not 0x1100 <= ord(text[0]) <= 0x1113:
        text = ''.join(list(hangul_to_jamo(text)))
    for s in text:
        sequence.append(_symbol_to_id[s])
    sequence.append(_symbol_to_id['~'])
    return sequence

def sequence_to_text(sequence):
    result = ''
    for symbol_id in sequence:
        if symbol_id in _id_to_symbol:
            s = _id_to_symbol[symbol_id]
            result += s
    return result.replace('}{', ' ')

def mel_spectrogram(y, n_fft=1024, num_mels=80, sampling_rate=22050, hop_size=256, win_size=1024, fmin=0, fmax=8000, center=False):
    """
    if torch.min(y) < -1.:
        print('min value is ', torch.min(y))
    if torch.max(y) > 1.:
        print('max value is ', torch.max(y))
    """

    mel = librosa.filters.mel(sampling_rate, n_fft, num_mels, fmin, fmax)
    
    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
    y = y.squeeze(1)

    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=torch.hann_window(win_size).to(y.device),
                      center=center, pad_mode='reflect', normalized=False, onesided=True)

    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))

    spec = torch.matmul(torch.from_numpy(mel).float().to(y.device), spec)
    spec = torch.log(torch.clamp(spec, min=1e-5) * 1)

    return spec