import numpy as np
import librosa
import glob
from hparams import hparams
import os

def quantize_data(data, classes):
    mu_x = mu_law_encoding(data, classes)
    bins = np.linspace(-1, 1, classes)
    quantized = np.digitize(mu_x, bins) - 1
    return quantized
    
def mu_law_encoding(data, mu):
    mu_x = np.sign(data) * np.log(1 + mu * np.abs(data)) / np.log(mu + 1)
    return mu_x


def mu_law_expansion(data, mu):
    s = np.sign(data) * (np.exp(np.abs(data) * np.log(mu + 1)) - 1) / mu
    return s    


if __name__ == "__main__":
    para = hparams()
    wavs = glob.glob(para.path_train_wavs+'/*wav') 
    os.makedirs(para.path_train_coded,exist_ok=True)
    
    # 遍历所有的wav 文件
    for file_wav in wavs:
        name = os.path.split(file_wav)[-1]
        audio,_ = librosa.load(file_wav,sr=para.fs,mono=True)
        # 进行mu 率量化编码
        quantized_data = quantize_data(audio, para.classes)
        save_name = os.path.join(para.path_train_coded,name+'.npy')
        np.save(save_name,quantized_data)