# From https://github.com/keunwoochoi/torchaudio-contrib/blob/master/torchaudio_contrib/functional.py

import torch
import math
import torch.nn.functional as F


def stft(signal, fft_len, hop_len, window,
         pad=0, pad_mode="reflect", **kwargs):
    """
    Wrap torch.stft allowing for multi-channel stft.

    Args:
        signal (Tensor): Tensor of audio of size (channel, time)
            or (batch, channel, time).
        fft_len (int): FFT window size.
        hop_len (int): Number audio of frames between STFT columns.
        window (Tensor): 1-D tensor.
        pad (int): Amount of padding to apply to signal.
        pad_mode: padding method (see torch.nn.functional.pad).
        **kwargs: Other torch.stft parameters, see torch.stft for more details.

    Returns:
        Tensor: (batch, channel, num_bins, time, complex)
            or (channel, num_bins, time, complex)

    Example:
        >>> signal = torch.randn(16, 2, 10000)
        >>> # window_length <= fft_len
        >>> window = torch.hamming_window(window_length=2048)
        >>> x = stft(signal, 2048, 512, window)
        >>> x.shape
        torch.Size([16, 2, 1025, 20])
    """

    # (!) Only 3D, 4D, 5D padding with non-constant
    # padding are supported for now.
    if pad > 0:
        signal = F.pad(signal, (pad, pad), pad_mode)

    leading_dims = signal.shape[:-1]

    signal = signal.reshape(-1, signal.size(-1))

    spect = torch.stft(signal, fft_len, hop_len, window=window,
                       win_length=window.size(0), **kwargs)
    spect = spect.reshape(leading_dims + spect.shape[1:])

    return spect


def complex_norm(tensor, power=1.0):
    """
    Normalize complex input.
    """
    if power == 1.:
        return torch.norm(tensor, 2, -1)
    return torch.norm(tensor, 2, -1).pow(power)


def create_mel_filter(num_bands, sample_rate, min_freq,
                      max_freq, num_bins, to_hertz, from_hertz):
    """
    Creates filter matrix to transform fft frequency bins
    into mel frequency bins.
    Equivalent to librosa.filters.mel(sample_rate, fft_len, htk=True, norm=None).

    Args:
        num_bands (int): number of mel bins.
        sample_rate (int): sample rate of audio signal.
        min_freq (float): minimum frequency.
        max_freq (float): maximum frequency.
        num_bins (int): number of filter banks from stft.
        to_hertz (function): convert from mel freq to hertz
        from_hertz (function): convert from hertz to mel freq

    Returns:
        filterbank (Tensor): (num_bins, num_bands)
    """
    # Convert to find mel lower/upper bounds
    m_min = from_hertz(min_freq)
    m_max = from_hertz(max_freq)

    # Compute stft frequency values
    stft_freqs = torch.linspace(min_freq, max_freq, num_bins)

    # Find mel values, and convert them to frequency units
    m_pts = torch.linspace(m_min, m_max, num_bands + 2)
    f_pts = to_hertz(m_pts)
    f_diff = f_pts[1:] - f_pts[:-1]  # (num_bands + 1)

    # (num_bins, num_bands + 2)
    slopes = f_pts.unsqueeze(0) - stft_freqs.unsqueeze(1)

    down_slopes = (-1. * slopes[:, :-2]) / f_diff[:-1]  # (num_bins, num_bands)
    up_slopes = slopes[:, 2:] / f_diff[1:]  # (num_bins, num_bands)
    filterbank = torch.clamp(torch.min(down_slopes, up_slopes), min=0.)

    return filterbank


def apply_filterbank(spect, filterbank):
    """
    Transform spectrogram given a filterbank matrix.

    Args:
        spect (Tensor): (batch, channel, num_bins, time)
        filterbank (Tensor): (num_bins, num_bands)

    Returns:
        (Tensor): (batch, channel, num_bands, time)
    """
    return torch.matmul(spect.transpose(-2, -1), filterbank).transpose(-2, -1)


def angle(tensor):
    """
    Return angle of a complex tensor with shape (*, 2).
    """
    return torch.atan2(tensor[..., 1], tensor[..., 0])


def magphase(spect, power=1.):
    """
    Separate a complex-valued spectrogram with shape (*,2)
    into its magnitude and phase.
    """
    mag = complex_norm(spect, power)
    phase = angle(spect)
    return mag, phase


def phase_vocoder(spect, rate, phi_advance):
    """
    Phase vocoder. Given a STFT tensor, speed up in time
    without modifying pitch by a factor of `rate`.

    Args:
        spect (Tensor): (batch, channel, num_bins, time, 2)
        rate (float): Speed-up factor
        phi_advance (Tensor): Expected phase advance in each bin. (num_bins, 1)

    Returns:
      (Tensor): (batch, channel, num_bins, new_bins, 2) with new_bins = num_bins//rate+1
    """

    time_steps = torch.arange(0, spect.size(
        3), rate, device=spect.device)  # (new_bins,)

    alphas = (time_steps % 1)  # (new_bins,)

    phase_0 = angle(spect[:, :, :, :1])

    # Time Padding
    pad_shape = [0, 0] + [0, 2] + [0] * 6
    spect = torch.nn.functional.pad(spect, pad_shape)

    spect_0 = spect[:, :, :, time_steps.long()]  # (new_bins, num_bins, 2)
    # (new_bins, num_bins, 2)
    spect_1 = spect[:, :, :, (time_steps + 1).long()]

    spect_0_angle = angle(spect_0)  # (new_bins, num_bins)
    spect_1_angle = angle(spect_1)  # (new_bins, num_bins)

    spect_0_norm = torch.norm(spect_0, dim=-1)  # (new_bins, num_bins)
    spect_1_norm = torch.norm(spect_1, dim=-1)  # (new_bins, num_bins)

    spect_phase = spect_1_angle - spect_0_angle - \
        phi_advance  # (new_bins, num_bins)
    spect_phase = spect_phase - 2 * math.pi * \
        torch.round(spect_phase / (2 * math.pi))  # (new_bins, num_bins)

    # Compute Phase Accum
    phase = spect_phase + phi_advance  # (new_bins, num_bins)

    phase = torch.cat([phase_0, phase[:, :, :, :-1]], dim=-1)

    phase_acc = torch.cumsum(phase, -1)  # (new_bins, num_bins)

    mag = alphas * spect_1_norm + (1 - alphas) * \
        spect_0_norm  # (time//rate+1, num_bins)

    spect_stretch_real = mag * torch.cos(phase_acc)  # (new_bins, num_bins)
    spect_stretch_imag = mag * torch.sin(phase_acc)  # (new_bins, num_bins)

    spect_stretch = torch.stack(
        [spect_stretch_real, spect_stretch_imag], dim=-1)

    return spect_stretch


def amplitude_to_db(x, ref=1.0, amin=1e-7):
    """
    Amplitude-to-decibel conversion (logarithmic mapping with base=10)
    By using `amin=1e-7`, it assumes 32-bit floating point input. If the
    data precision differs, use approproate `amin` accordingly.

    Args:
        x (Tensor): Input amplitude
        ref (float): Amplitude value that is equivalent to 0 decibel
        amin (float): Minimum amplitude. Any input that is smaller than `amin` is
            clamped to `amin`.
    Returns:
        (Tensor): same size of x, after conversion
    """
    x = torch.clamp(x, min=amin)
    return 10.0 * (torch.log10(x) - torch.log10(torch.tensor(ref, device=x.device, requires_grad=False)))


def db_to_amplitude(x, ref=1.0):
    """
    Decibel-to-amplitude conversion (exponential mapping with base=10)

    Args:
        x (Tensor): Input in decibel to be converted
        ref (float): Amplitude value that is equivalent to 0 decibel

    Returns:
        (Tensor): same size of x, after conversion
    """
    return torch.pow(10.0, x / 10.0 + torch.log10(torch.tensor(ref, device=x.device, requires_grad=False)))


def mu_law_encoding(x, n_quantize=256):
    """Apply mu-law encoding to the input tensor.
    Usually applied to waveforms

    Args:
        x (Tensor): input value
        n_quantize (int): quantization level. For 8-bit encoding, set 256 (2 ** 8).

    Returns:
        (Tensor): same size of x, after encoding

    """
    if not x.dtype.is_floating_point:
        x = x.to(torch.float)
    mu = torch.tensor(n_quantize - 1, dtype=x.dtype, requires_grad=False)  # confused about dtype here..

    x_mu = x.sign() * torch.log1p(mu * x.abs()) / torch.log1p(mu)
    x_mu = ((x_mu + 1) / 2 * mu + 0.5).long()
    return x_mu


def mu_law_decoding(x_mu, n_quantize=256):
    """Apply mu-law decoding (expansion) to the input tensor.

    Args:
        x_mu (Tensor): mu-law encoded input
        n_quantize (int): quantization level. For 8-bit decoding, set 256 (2 ** 8).

    Returns:
        (Tensor): mu-law decoded tensor
    """
    if not x_mu.dtype.is_floating_point:
        x_mu = x_mu.to(torch.float)
    mu = torch.tensor(n_quantize - 1, dtype=x_mu.dtype, requires_grad=False)  # confused about dtype here..
    x = (x_mu / mu) * 2 - 1.
    x = x.sign() * (torch.exp(x.abs() * torch.log1p(mu)) - 1.) / mu
    return x

# From https://github.com/keunwoochoi/torchaudio-contrib/blob/master/torchaudio_contrib/layers.py
class Filterbank(object):
    """
    Base class for providing a filterbank matrix.
    """

    def __init__(self):
        super(Filterbank, self).__init__()

    def get_filterbank(self):
        raise NotImplementedError

class MelFilterbank(Filterbank):
    """
    Provides a filterbank matrix to convert a spectrogram into a mel frequency spectrogram.
    Args:
        num_bands (int): number of mel bins. Defaults to 128.
        sample_rate (int): sample rate of audio signal. Defaults to 22050.
        min_freq (float): minimum frequency. Defaults to 0.
        max_freq (float, optional): maximum frequency. Defaults to sample_rate // 2.
        num_bins (int, optional): number of filter banks from stft.
            Defaults to 2048//2 + 1.
        htk (bool, optional): use HTK formula instead of Slaney. Defaults to False.
    """

    def __init__(self, num_bands=128, sample_rate=22050,
                 min_freq=0.0, max_freq=None, num_bins=1025, htk=False):

        super(MelFilterbank, self).__init__()

        self.num_bands = num_bands
        self.sample_rate = sample_rate
        self.min_freq = min_freq
        self.max_freq = max_freq if max_freq else sample_rate // 2
        self.num_bins = num_bins
        self.htk = htk

    def to_hertz(self, mel):
        """
        Converting mel values into frequency
        """
        mel = torch.as_tensor(mel).float()

        if self.htk:
            return 700. * (10**(mel / 2595.) - 1.)

        f_min = 0.0
        f_sp = 200.0 / 3
        hz = f_min + f_sp * mel

        min_log_hz = 1000.0
        min_log_mel = (min_log_hz - f_min) / f_sp
        logstep = math.log(6.4) / 27.0

        return torch.where(mel >= min_log_mel, min_log_hz *
                           torch.exp(logstep * (mel - min_log_mel)), hz)

    def from_hertz(self, hz):
        """
        Converting frequency into mel values
        """
        hz = torch.as_tensor(hz).float()

        if self.htk:
            return 2595. * torch.log10(torch.tensor(1.) + (hz / 700.))

        f_min = 0.0
        f_sp = 200.0 / 3

        mel = (hz - f_min) / f_sp

        min_log_hz = 1000.0
        min_log_mel = (min_log_hz - f_min) / f_sp
        logstep = math.log(6.4) / 27.0

        return torch.where(hz >= min_log_hz, min_log_mel +
                           torch.log(hz / min_log_hz) / logstep, mel)

    def get_filterbank(self):
        return create_mel_filter(
            num_bands=self.num_bands,
            sample_rate=self.sample_rate,
            min_freq=self.min_freq,
            max_freq=self.max_freq,
            num_bins=self.num_bins,
            to_hertz=self.to_hertz,
            from_hertz=self.from_hertz)

    def __repr__(self):
        param_str1 = '(num_bands={}, sample_rate={}'.format(
            self.num_bands, self.sample_rate)
        param_str2 = ', min_freq={}, max_freq={})'.format(
            self.min_freq, self.max_freq)
        return self.__class__.__name__ + param_str1 + param_str2

