import math
import os
import random
import torch
import torch.utils.data
import numpy as np
import pandas as pd
from librosa.util import normalize
from scipy.io.wavfile import read
import librosa
from librosa.filters import mel as librosa_mel_fn
import audio as Audio
from audio import stft

MAX_WAV_VALUE = 32768.0


def pcm2float(sig, dtype='float32'):
    """Convert PCM signal to floating point with a range from -1 to 1.
    Use dtype='float32' for single precision.
    Parameters
    ----------
    sig : array_like
        Input array, must have integral type.
    dtype : data type, optional
        Desired (floating point) data type.
    Returns
    -------
    numpy.ndarray
        Normalized floating point data.
    See Also
    --------
    float2pcm, dtype
    """
    sig = np.asarray(sig)
    if sig.dtype.kind not in 'iu':
        raise TypeError("'sig' must be an array of integers")
    dtype = np.dtype(dtype)
    if dtype.kind != 'f':
        raise TypeError("'dtype' must be a floating point type")

    i = np.iinfo(sig.dtype)
    abs_max = 2 ** (i.bits - 1)
    offset = i.min + abs_max
    return (sig.astype(dtype) - offset) / abs_max

def load_wav(full_path):
    sampling_rate, data = read(full_path)
    return data, sampling_rate


def dynamic_range_compression(x, C=1, clip_val=1e-5):
    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)


def dynamic_range_decompression(x, C=1):
    return np.exp(x) / C


def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
    return torch.log(torch.clamp(x, min=clip_val) * C)


def dynamic_range_decompression_torch(x, C=1):
    return torch.exp(x) / C


def spectral_normalize_torch(magnitudes):
    output = dynamic_range_compression_torch(magnitudes)
    return output


def spectral_de_normalize_torch(magnitudes):
    output = dynamic_range_decompression_torch(magnitudes)
    return output


mel_basis = {}
hann_window = {}


def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, 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))

    global mel_basis, hann_window
    if fmax not in mel_basis:
        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
        mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)

    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft)/2), int((n_fft)/2)), mode='reflect')
    y = y.squeeze(1)

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

    #print(spec.shape)
    #print(spec)
    #import ipdb; ipdb.set_trace()
    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
    print(spec.squeeze(0)[0], spec.shape)
    print(spec.shape)

    spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec)
    spec = spectral_normalize_torch(spec)

    return spec

if __name__ == '__main__':

    
    _path='/data/voice_team/cooked_data/voice_conversion/voice_actor_pubg/audio/김민주/거짓말퍼즐/E10_0837.wav'
    np_audio, rate = librosa.load(_path, sr=16000)
    audio = torch.FloatTensor(np_audio)
    audio = audio.unsqueeze(0)

    ref_mel = np.load('/data/voice_team/cooked_data/voice_conversion/voice_actor_pubg/mel/uid0_E10_0837.npy')

    mel = mel_spectrogram(audio, n_fft=1024, num_mels=80, sampling_rate=16000,
                            hop_size=256, win_size=1024, fmin=0, fmax=8000, center=False)
    lib_mel = np.log(np.clip(librosa.feature.melspectrogram(np_audio, n_mels=80, sr=16000, n_fft=1024,
            hop_length=256,win_length=1024,power=1.0), a_min=1e-5,a_max=np.inf))

    _stft = stft.TacotronSTFT(
        1024, 256, 1024,
        80, 16000, 0, 80)
    mel_spectogram, energy = Audio.tools.get_mel_from_wav(
            torch.FloatTensor(np_audio), _stft, 1.0)


#    print(mel.squeeze(0).T)
#    print('========================')
#    print(ref_mel)
#    print('========================')
#    print(lib_mel.T)
#    print('========================')
#    print(mel_spectrogram)
#    print('========================')
