import typing
from typing import List

import librosa
import torch
import torch.nn.functional as F
import numpy as np
from einops import rearrange
from audiotools import AudioSignal
from audiotools import STFTParams
from torch import nn


def PQWOME(f):
    # Generate the weighting for the outer & middle ear filtering
    # Note: The output is a magnitude-squared vector

    N = len(f)
    W2 = np.zeros(N)

    for k in range(N - 1):
        fkHz = float(f[k + 1]) / 1000
        AdB = (
            -2.184 * fkHz ** (-0.8)
            + 6.5 * np.exp(-0.6 * (fkHz - 3.3) ** 2)
            - 0.001 * fkHz ** (3.6)
        )
        W2[k + 1] = 10 ** (AdB / 10)

    return W2


def safe_pow(x, y):
    if x == 0 and y <= 0:
        if y == 0:
            return 1
        else:
            return float("inf")
    return x**y


def omem(f):
    return float(
        -2.184 * safe_pow(f / 1000, -0.8)
        + 6.5 * np.exp(-0.6 * safe_pow(f / 1000 - 3.3, 2))
        - 0.001 * safe_pow(f / 1000, 3.6)
    )


class L1Loss(nn.L1Loss):
    """L1 Loss between AudioSignals. Defaults
    to comparing ``audio_data``, but any
    attribute of an AudioSignal can be used.

    Parameters
    ----------
    attribute : str, optional
        Attribute of signal to compare, defaults to ``audio_data``.
    weight : float, optional
        Weight of this loss, defaults to 1.0.

    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
    """

    def __init__(self, attribute: str = "audio_data", weight: float = 1.0, **kwargs):
        self.attribute = attribute
        self.weight = weight
        super().__init__(**kwargs)

    def forward(self, x: AudioSignal, y: AudioSignal):
        """
        Parameters
        ----------
        x : AudioSignal
            Estimate AudioSignal
        y : AudioSignal
            Reference AudioSignal

        Returns
        -------
        torch.Tensor
            L1 loss between AudioSignal attributes.
        """
        if isinstance(x, AudioSignal):
            x = getattr(x, self.attribute)
            y = getattr(y, self.attribute)
        return super().forward(x, y)


class SISDRLoss(nn.Module):
    """
    Computes the Scale-Invariant Source-to-Distortion Ratio between a batch
    of estimated and reference audio signals or aligned features.

    Parameters
    ----------
    scaling : int, optional
        Whether to use scale-invariant (True) or
        signal-to-noise ratio (False), by default True
    reduction : str, optional
        How to reduce across the batch (either 'mean',
        'sum', or none).], by default ' mean'
    zero_mean : int, optional
        Zero mean the references and estimates before
        computing the loss, by default True
    clip_min : int, optional
        The minimum possible loss value. Helps network
        to not focus on making already good examples better, by default None
    weight : float, optional
        Weight of this loss, defaults to 1.0.

    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
    """

    def __init__(
        self,
        scaling: int = True,
        reduction: str = "mean",
        zero_mean: int = True,
        clip_min: int = None,
        weight: float = 1.0,
    ):
        self.scaling = scaling
        self.reduction = reduction
        self.zero_mean = zero_mean
        self.clip_min = clip_min
        self.weight = weight
        super().__init__()

    def forward(self, x: AudioSignal, y: AudioSignal):
        eps = 1e-8
        # nb, nc, nt
        if isinstance(x, AudioSignal):
            references = x.audio_data
            estimates = y.audio_data
        else:
            references = x
            estimates = y

        nb = references.shape[0]
        references = references.reshape(nb, 1, -1).permute(0, 2, 1)
        estimates = estimates.reshape(nb, 1, -1).permute(0, 2, 1)

        # samples now on axis 1
        if self.zero_mean:
            mean_reference = references.mean(dim=1, keepdim=True)
            mean_estimate = estimates.mean(dim=1, keepdim=True)
        else:
            mean_reference = 0
            mean_estimate = 0

        _references = references - mean_reference
        _estimates = estimates - mean_estimate

        references_projection = (_references**2).sum(dim=-2) + eps
        references_on_estimates = (_estimates * _references).sum(dim=-2) + eps

        scale = (
            (references_on_estimates / references_projection).unsqueeze(1)
            if self.scaling
            else 1
        )

        e_true = scale * _references
        e_res = _estimates - e_true

        signal = (e_true**2).sum(dim=1)
        noise = (e_res**2).sum(dim=1)
        sdr = -10 * torch.log10(signal / noise + eps)

        if self.clip_min is not None:
            sdr = torch.clamp(sdr, min=self.clip_min)

        if self.reduction == "mean":
            sdr = sdr.mean()
        elif self.reduction == "sum":
            sdr = sdr.sum()
        return sdr


class MultiScaleSTFTLoss(nn.Module):
    """Computes the multi-scale STFT loss from [1].

    Parameters
    ----------
    window_lengths : List[int], optional
        Length of each window of each STFT, by default [2048, 512]
    loss_fn : typing.Callable, optional
        How to compare each loss, by default nn.L1Loss()
    clamp_eps : float, optional
        Clamp on the log magnitude, below, by default 1e-5
    mag_weight : float, optional
        Weight of raw magnitude portion of loss, by default 1.0
    log_weight : float, optional
        Weight of log magnitude portion of loss, by default 1.0
    pow : float, optional
        Power to raise magnitude to before taking log, by default 2.0
    weight : float, optional
        Weight of this loss, by default 1.0
    match_stride : bool, optional
        Whether to match the stride of convolutional layers, by default False

    References
    ----------

    1.  Engel, Jesse, Chenjie Gu, and Adam Roberts.
        "DDSP: Differentiable Digital Signal Processing."
        International Conference on Learning Representations. 2019.

    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
    """

    def __init__(
        self,
        window_lengths: List[int] = [2048, 512],
        loss_fn: typing.Callable = nn.L1Loss(),
        clamp_eps: float = 1e-5,
        mag_weight: float = 1.0,
        log_weight: float = 1.0,
        pow: float = 2.0,
        weight: float = 1.0,
        match_stride: bool = False,
        window_type: str = None,
        is_perceptual: bool = False,
    ):
        super().__init__()
        self.stft_params = [
            STFTParams(
                window_length=w,
                hop_length=w // 4,
                match_stride=match_stride,
                window_type=window_type,
            )
            for w in window_lengths
        ]
        self.loss_fn = loss_fn
        self.log_weight = log_weight
        self.mag_weight = mag_weight
        self.clamp_eps = clamp_eps
        self.weight = weight
        self.pow = pow
        self.is_perceptual = is_perceptual

    def forward(self, x: AudioSignal, y: AudioSignal):
        """Computes multi-scale STFT between an estimate and a reference
        signal.

        Parameters
        ----------
        x : AudioSignal
            Estimate signal
        y : AudioSignal
            Reference signal

        Returns
        -------
        torch.Tensor
            Multi-scale STFT loss.
        """
        loss = 0.0
        for s in self.stft_params:
            x.stft(s.window_length, s.hop_length, s.window_type)
            y.stft(s.window_length, s.hop_length, s.window_type)
            if self.is_perceptual:
                f = np.linspace(0, x.sample_rate / 2, s.window_length / 2 + 1)
                ### Check hsape
                pq = PQWOME(f)[np.newaxis, np.newaxis, :]
                loss += self.log_weight * self.loss_fn(
                    x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
                    y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
                )
                loss += self.log_weight * self.loss_fn(
                    torch.view_as_real(x.stft_data) / 60.0,
                    torch.view_as_real(y.stft_data) / 60.0,
                )
                loss += self.mag_weight * self.loss_fn(x.magnitude, y.magnitude)
            else:
                loss += self.log_weight * self.loss_fn(
                    x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
                    y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
                )
                loss += self.log_weight * self.loss_fn(
                    torch.view_as_real(x.stft_data) / 60.0,
                    torch.view_as_real(y.stft_data) / 60.0,
                )
                loss += self.mag_weight * self.loss_fn(x.magnitude, y.magnitude)
        return loss


class MelSpectrogramLoss(nn.Module):
    """Compute distance between mel spectrograms. Can be used
    in a multi-scale way.

    Parameters
    ----------
    n_mels : List[int]
        Number of mels per STFT, by default [150, 80],
    window_lengths : List[int], optional
        Length of each window of each STFT, by default [2048, 512]
    loss_fn : typing.Callable, optional
        How to compare each loss, by default nn.L1Loss()
    clamp_eps : float, optional
        Clamp on the log magnitude, below, by default 1e-5
    mag_weight : float, optional
        Weight of raw magnitude portion of loss, by default 1.0
    log_weight : float, optional
        Weight of log magnitude portion of loss, by default 1.0
    pow : float, optional
        Power to raise magnitude to before taking log, by default 2.0
    weight : float, optional
        Weight of this loss, by default 1.0
    match_stride : bool, optional
        Whether to match the stride of convolutional layers, by default False

    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
    """

    def __init__(
        self,
        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
        loss_fn: typing.Callable = nn.L1Loss(),
        clamp_eps: float = 1e-5,
        mag_weight: float = 1.0,
        log_weight: float = 1.0,
        pow: float = 2.0,
        weight: float = 1.0,
        match_stride: bool = False,
        mel_fmin: List[float] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        mel_fmax: List[float] = [None, None, None, None, None, None, None],
        window_type: str = None,
    ):
        super().__init__()
        self.stft_params = [
            STFTParams(
                window_length=w,
                hop_length=w // 4,
                match_stride=match_stride,
                window_type=window_type,
            )
            for w in window_lengths
        ]
        self.n_mels = n_mels
        self.loss_fn = loss_fn
        self.clamp_eps = clamp_eps
        self.log_weight = log_weight
        self.mag_weight = mag_weight
        self.weight = weight
        self.mel_fmin = mel_fmin
        self.mel_fmax = mel_fmax
        self.pow = pow

    def forward(self, x: AudioSignal, y: AudioSignal):
        """Computes mel loss between an estimate and a reference
        signal.

        Parameters
        ----------
        x : AudioSignal
            Estimate signal
        y : AudioSignal
            Reference signal

        Returns
        -------
        torch.Tensor
            Mel loss.
        """
        loss = 0.0
        for n_mels, fmin, fmax, s in zip(
            self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params
        ):
            kwargs = {
                "window_length": s.window_length,
                "hop_length": s.hop_length,
                "window_type": s.window_type,
            }
            x_mels = x.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)
            y_mels = y.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)

            loss += self.log_weight * self.loss_fn(
                x_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
                y_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
            )
            loss += self.mag_weight * self.loss_fn(x_mels, y_mels)
        return loss


class GANLoss(nn.Module):
    """
    Computes a discriminator loss, given a discriminator on
    generated waveforms/spectrograms compared to ground truth
    waveforms/spectrograms. Computes the loss for both the
    discriminator and the generator in separate functions.
    """

    def __init__(self, discriminator):
        super().__init__()
        self.discriminator = discriminator

    def forward(self, fake, real):
        d_fake = self.discriminator(fake.audio_data)
        d_real = self.discriminator(real.audio_data)
        return d_fake, d_real

    def discriminator_loss(self, fake, real):
        d_fake, d_real = self.forward(fake.clone().detach(), real)

        loss_d = 0
        for x_fake, x_real in zip(d_fake, d_real):
            loss_d += torch.mean(x_fake[-1] ** 2)
            loss_d += torch.mean((1 - x_real[-1]) ** 2)
        return loss_d

    def generator_loss(self, fake, real):
        d_fake, d_real = self.forward(fake, real)

        loss_g = 0
        for x_fake in d_fake:
            loss_g += torch.mean((1 - x_fake[-1]) ** 2)

        loss_feature = 0

        for i in range(len(d_fake)):
            for j in range(len(d_fake[i]) - 1):
                loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())
        return loss_g, loss_feature


class BandSplitSpectrogramLoss(nn.Module):
    """Compute distance between band-split spectrograms.

    Parameters
    ----------
    n_bands : Lint[int], optional
        Number of split bands
    window_lengths : List[int], optional
        Length of each window of each STFT, by default [2048, 512]
    loss_fn : typing.Callable, optional
        How to compare each loss, by default nn.L1Loss()
    clamp_eps : float, optional
        Clamp on the log magnitude, below, by default 1e-5
    mag_weight : float, optional
        Weight of raw magnitude portion of loss, by default 1.0
    log_weight : float, optional
        Weight of log magnitude portion of loss, by default 1.0
    pow : float, optional
        Power to raise magnitude to before taking log, by default 2.0
    weight : float, optional
        Weight of this loss, by default 1.0
    match_stride : bool, optional
        Whether to match the stride of convolutional layers, by default False

    Implemented by Minz Won
    """

    def __init__(
        self,
        sample_rate: int = 48000,
        n_bands: List[int] = [4, 8, 16, 32, 64, 128, 256],
        window_lengths: List[int] = [32, 64, 128, 512, 1024, 2048, 4096],
        loss_fn: typing.Callable = nn.L1Loss(reduction="none"),
        clamp_eps: float = 1e-5,
        log_weight: float = 1.0,
        pow: float = 2.0,
        loss_weight: float = 1.0,
        match_stride: bool = False,
        window_type: str = None,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.n_bands = n_bands
        self.stft_params = [
            STFTParams(
                window_length=w,
                hop_length=w // 4,
                match_stride=match_stride,
                window_type=window_type,
            )
            for w in window_lengths
        ]
        self.loss_fn = loss_fn
        self.log_weight = log_weight
        self.clamp_eps = clamp_eps
        self.loss_weight = loss_weight
        self.pow = pow
        self.band_indices, self.band_kernels = self.preload_band_indices()

    def get_band_indices(self, n_fft, n_band):
        # get band indices
        mel_basis = librosa.filters.mel(sr=self.sample_rate, n_fft=n_fft, n_mels=n_band)
        band_indices = [np.where(row > 0)[0] for row in mel_basis]

        # make kernels
        kernels = (mel_basis > 0).astype("float32")
        kernels = kernels / kernels.sum(axis=1)[:, np.newaxis]
        kernels = torch.tensor(kernels).T
        return band_indices, kernels

    def preload_band_indices(self):
        band_indices = {}
        kernels = {}
        for s in self.stft_params:
            for b in self.n_bands:
                if s.window_length / b >= 8:
                    # get band indices
                    (
                        band_indices[f"{s.window_length}_{b}"],
                        kernels[f"{s.window_length}_{b}"],
                    ) = self.get_band_indices(s.window_length, b)
        return band_indices, kernels

    def forward(self, x: AudioSignal, y: AudioSignal):
        """Computes multi-scale STFT between an estimate and a reference
        signal.

        Parameters
        ----------
        x : AudioSignal [batch * 2, 1, time]
            Estimate signal in stereo
        y : AudioSignal [batch * 2, 1, time]
            Reference signal in stereo

        Returns
        -------
        torch.Tensor
            Multi-scale Band-split spectrogram loss
        """
        loss = 0.0
        cnt = 0
        for s in self.stft_params:  # for multiple window lengths
            x.stft(s.window_length, s.hop_length, s.window_type)
            y.stft(s.window_length, s.hop_length, s.window_type)
            stft_loss = self.loss_fn(
                x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
                y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
            )
            for b in self.n_bands:  # for multiple band splits
                if s.window_length / b >= 8:  # discard when lower band is empty
                    loss += torch.einsum(
                        "bcft, fn -> bcnt",
                        stft_loss,
                        self.band_kernels[f"{s.window_length}_{b}"].to(x.device),
                    ).mean()
                    cnt += 1
        loss = self.loss_weight * loss / cnt

        return loss


class PEAQL1Loss(nn.Module):
    """Compute distance between Bark-scale spectrograms described in PEAQ.

    Parameters
    ----------

    """

    def __init__(
        self,
        sample_rate: int = 48000,
        window_length: int = 2048,
        overlap: int = 1024,
        loss_fn: typing.Callable = nn.L1Loss(),
    ):
        super().__init__()

        # hyper parameters
        self.sample_rate = sample_rate
        self.window_length = window_length
        self.overlap = overlap
        self.loss_fn = loss_fn

        # fixed parameters
        L_p = 92
        A_max = 1
        gamma_fc = 0.8497
        G_L = 10 ** (L_p / 20) / (gamma_fc * A_max / 4 * (window_length - 1))
        hann = torch.hann_window(window_length).numpy()
        weighted_hann = G_L * hann
        self.weighted_hann = weighted_hann[np.newaxis, :]
        self.omem_weights = self.get_omem()
        self.filterbank, self.f_c = self.get_filterbank()
        self.internal_noise = self.get_noise()

    def get_omem(self):
        """Outher and middle ear modeling weights"""
        center_freq = np.linspace(
            0, self.sample_rate // 2, self.window_length // 2 + 1
        ).astype(np.float32)
        a_db = [omem(f) for f in center_freq]
        weights = torch.tensor([10 ** (a / 20) for a in a_db]).unsqueeze(0)
        return weights

    def get_filterbank(self):
        """Bark-scale filterbank"""
        f_l = np.array(
            [
                80.000,
                103.445,
                127.023,
                150.762,
                174.694,
                198.849,
                223.257,
                247.950,
                272.959,
                298.317,
                324.055,
                350.207,
                376.805,
                403.884,
                431.478,
                459.622,
                488.353,
                517.707,
                547.721,
                578.434,
                609.885,
                642.114,
                675.161,
                709.071,
                743.884,
                779.647,
                816.404,
                854.203,
                893.091,
                933.119,
                974.336,
                1016.797,
                1060.555,
                1105.666,
                1152.187,
                1200.178,
                1249.700,
                1300.816,
                1353.592,
                1408.094,
                1464.392,
                1522.559,
                1582.668,
                1644.795,
                1709.021,
                1775.427,
                1844.098,
                1915.121,
                1988.587,
                2064.590,
                2143.227,
                2224.597,
                2308.806,
                2395.959,
                2486.169,
                2579.551,
                2676.223,
                2776.309,
                2879.937,
                2987.238,
                3098.350,
                3213.415,
                3332.579,
                3455.993,
                3583.817,
                3716.212,
                3853.817,
                3995.399,
                4142.547,
                4294.979,
                4452.890,
                4616.482,
                4785.962,
                4961.548,
                5143.463,
                5331.939,
                5527.217,
                5729.545,
                5939.183,
                6156.396,
                6381.463,
                6614.671,
                6856.316,
                7106.708,
                7366.166,
                7635.020,
                7913.614,
                8202.302,
                8501.454,
                8811.450,
                9132.688,
                9465.574,
                9810.536,
                10168.013,
                10538.460,
                10922.351,
                11320.175,
                11732.438,
                12159.670,
                12602.412,
                13061.229,
                13536.710,
                14029.458,
                14540.103,
                15069.295,
                15617.710,
                16186.049,
                16775.035,
                17385.420,
            ]
        )
        f_c = np.array(
            [
                91.708,
                115.216,
                138.870,
                162.702,
                186.742,
                211.019,
                235.566,
                260.413,
                285.593,
                311.136,
                337.077,
                363.448,
                390.282,
                417.614,
                445.479,
                473.912,
                502.950,
                532.629,
                562.988,
                594.065,
                625.899,
                658.533,
                692.006,
                726.362,
                761.644,
                797.898,
                835.170,
                873.508,
                912.959,
                953.576,
                995.408,
                1038.511,
                1082.938,
                1128.746,
                1175.995,
                1224.744,
                1275.055,
                1326.992,
                1380.623,
                1436.014,
                1493.237,
                1552.366,
                1613.474,
                1676.641,
                1741.946,
                1809.474,
                1879.310,
                1951.543,
                2026.266,
                2103.573,
                2183.564,
                2266.340,
                2352.008,
                2440.675,
                2532.456,
                2627.468,
                2725.832,
                2827.672,
                2933.120,
                3042.309,
                3155.379,
                3272.475,
                3393.745,
                3519.344,
                3649.432,
                3784.176,
                3923.748,
                4068.324,
                4218.090,
                4373.237,
                4533.963,
                4700.473,
                4872.978,
                5051.700,
                5236.866,
                5428.712,
                5627.484,
                5833.434,
                6046.825,
                6267.931,
                6497.031,
                6734.420,
                6980.399,
                7235.284,
                7499.397,
                7773.077,
                8056.673,
                8350.547,
                8655.072,
                8970.639,
                9297.648,
                9636.520,
                9987.683,
                10351.586,
                10728.695,
                11119.490,
                11524.470,
                11944.149,
                12379.066,
                12829.775,
                13294.850,
                13780.887,
                14282.503,
                14802.338,
                15341.057,
                15899.345,
                16477.914,
                17077.504,
                17690.045,
            ]
        )
        f_u = np.array(
            [
                103.445,
                127.023,
                150.762,
                174.694,
                198.849,
                223.257,
                247.950,
                272.959,
                298.317,
                324.055,
                350.207,
                376.805,
                403.884,
                431.478,
                459.622,
                488.353,
                517.707,
                547.721,
                578.434,
                609.885,
                642.114,
                675.161,
                709.071,
                743.884,
                779.647,
                816.404,
                854.203,
                893.091,
                933.113,
                974.336,
                1016.797,
                1060.555,
                1105.666,
                1152.187,
                1200.178,
                1249.700,
                1300.816,
                1353.592,
                1408.094,
                1464.392,
                1522.559,
                1582.668,
                1644.795,
                1709.021,
                1775.427,
                1844.098,
                1915.121,
                1988.587,
                2064.590,
                2143.227,
                2224.597,
                2308.806,
                2395.959,
                2486.169,
                2579.551,
                2676.223,
                2776.309,
                2879.937,
                2987.238,
                3098.350,
                3213.415,
                3332.579,
                3455.993,
                3583.817,
                3716.212,
                3853.348,
                3995.399,
                4142.547,
                4294.979,
                4452.890,
                4643.482,
                4785.962,
                4961.548,
                5143.463,
                5331.939,
                5527.217,
                5729.545,
                5939.183,
                6156.396,
                6381.463,
                6614.671,
                6856.316,
                7106.708,
                7366.166,
                7635.020,
                7913.614,
                8202.302,
                8501.454,
                8811.450,
                9132.688,
                9465.574,
                9810.536,
                10168.013,
                10538.460,
                10922.351,
                11320.175,
                11732.438,
                12159.670,
                12602.412,
                13061.229,
                13536.710,
                14029.458,
                14540.103,
                15069.295,
                15617.710,
                16186.049,
                16775.035,
                17385.420,
                18000.000,
            ]
        )
        df = self.sample_rate / self.window_length
        filterbank = torch.zeros(self.window_length // 2 + 1, len(f_c))
        for k in range(filterbank.shape[0]):
            for i in range(filterbank.shape[1]):
                temp = (
                    np.amin([f_u[i], (k + 0.5) * df])
                    - np.amax([f_l[i], (k - 0.5) * df])
                ) / df
                filterbank[k, i] = np.amax([0, temp])
        return filterbank, f_c

    def get_noise(self):
        """Internal ear noise"""
        return torch.tensor(
            [10 ** (1.456 * ((f / 1000) ** -0.8) / 10) for f in self.f_c]
        ).unsqueeze(0)

    def slicing(self, sequence: np.array):
        """
        Input
            sequence: numpy array with a shape [batch, length]
        Output
            numpy array with a shape [batch, num_chunks, window_length]
        """
        num_chunks = (sequence.shape[-1] - self.window_length) // (
            self.window_length - self.overlap
        ) + 1
        chunks = np.stack(
            [
                np.lib.stride_tricks.sliding_window_view(s, self.window_length)[
                    :: self.window_length - self.overlap
                ]
                for s in sequence
            ]
        )  # batch, num_slice, window_length
        return chunks[:, :num_chunks, :]

    def windowing(self, stacked_wav):
        """Hann window"""
        return stacked_wav * self.weighted_hann

    def DFT(self, windowed_wav):
        """Discrete Fourier transform"""
        effective_bins = self.window_length // 2
        spec = torch.fft.fft(torch.tensor(windowed_wav), n=self.window_length)
        squared_spec = torch.cat(
            [
                spec[:, :1].real ** 2,
                spec[:, 1:effective_bins].real ** 2
                + spec[:, 1:effective_bins].imag ** 2,
                spec[:, effective_bins : effective_bins + 1].real ** 2,
            ],
            dim=-1,
        )
        return squared_spec

    def get_representation(self, audio: AudioSignal):
        """
        Input
            audio: AudioSignal
        """
        # input audio
        wav = audio.audio_data.squeeze(1).cpu().detach().numpy()
        b = len(wav)

        # slicing
        stacked_wav = self.slicing(wav)
        stacked_wav = rearrange(stacked_wav, "b c w -> (b c) w")

        # windowing
        windowed_wav = self.windowing(stacked_wav)

        # discrete Fourier transform
        squared_spec = self.DFT(windowed_wav)

        # outer and middle ear modeling
        weighted_spec = (self.omem_weights**2) * squared_spec

        # grouping filterbank
        grouped_spec = torch.matmul(weighted_spec, self.filterbank)
        grouped_spec[grouped_spec < 1e-12] = 1e-12

        # internal ear noise
        noised_spec = grouped_spec + self.internal_noise

        # rearrange
        return rearrange(noised_spec, "(b c) f -> b c f", b=b)

    def forward(self, x: AudioSignal, y: AudioSignal):
        """Computes mel loss between an estimate and a reference
        signal.

        Parameters
        ----------
        x : AudioSignal
            Estimate signal
        y : AudioSignal
            Reference signal

        Returns
        -------
        torch.Tensor
            Mel loss.
        """
        x_repr = self.get_representation(x)
        y_repr = self.get_representation(y)

        x_spec = 20 * torch.log10(x_repr)
        y_spec = 20 * torch.log10(y_repr)

        loss = self.loss_fn(x_spec, y_spec)
        return loss


class PEAQSpreadL1Loss(nn.Module):
    """Compute distance between spreaded Bark-scale spectrograms described in PEAQ.

    Parameters
    ----------

    """

    def __init__(
        self,
        sample_rate: int = 48000,
        window_length: int = 2048,
        overlap: int = 1024,
        loss_fn: typing.Callable = nn.L1Loss(),
        ratio: float = 0.8,
    ):
        super().__init__()

        # hyper parameters
        self.sample_rate = sample_rate
        self.window_length = window_length
        self.overlap = overlap
        self.loss_fn = loss_fn
        self.ratio = ratio

        # fixed parameters
        L_p = 92
        A_max = 1
        gamma_fc = 0.8497
        G_L = 10 ** (L_p / 20) / (gamma_fc * A_max / 4 * (window_length - 1))
        hann = torch.hann_window(window_length).numpy()
        weighted_hann = G_L * hann
        self.weighted_hann = weighted_hann[np.newaxis, :]
        self.omem_weights = self.get_omem()
        self.filterbank, self.f_c = self.get_filterbank()
        self.internal_noise = self.get_noise()
        self.bs = self.frequency_spreading(
            torch.ones(1, len(self.f_c)), torch.ones(1, len(self.f_c))
        )
        self.alpha = self.get_alpha()

    def get_omem(self):
        """Outher and middle ear modeling weights"""
        center_freq = np.linspace(
            0, self.sample_rate // 2, self.window_length // 2 + 1
        ).astype(np.float32)
        a_db = [omem(f) for f in center_freq]
        weights = torch.tensor([10 ** (a / 20) for a in a_db]).unsqueeze(0)
        return weights

    def get_filterbank(self):
        """Bark-scale filterbank"""
        f_l = np.array(
            [
                80.000,
                103.445,
                127.023,
                150.762,
                174.694,
                198.849,
                223.257,
                247.950,
                272.959,
                298.317,
                324.055,
                350.207,
                376.805,
                403.884,
                431.478,
                459.622,
                488.353,
                517.707,
                547.721,
                578.434,
                609.885,
                642.114,
                675.161,
                709.071,
                743.884,
                779.647,
                816.404,
                854.203,
                893.091,
                933.119,
                974.336,
                1016.797,
                1060.555,
                1105.666,
                1152.187,
                1200.178,
                1249.700,
                1300.816,
                1353.592,
                1408.094,
                1464.392,
                1522.559,
                1582.668,
                1644.795,
                1709.021,
                1775.427,
                1844.098,
                1915.121,
                1988.587,
                2064.590,
                2143.227,
                2224.597,
                2308.806,
                2395.959,
                2486.169,
                2579.551,
                2676.223,
                2776.309,
                2879.937,
                2987.238,
                3098.350,
                3213.415,
                3332.579,
                3455.993,
                3583.817,
                3716.212,
                3853.817,
                3995.399,
                4142.547,
                4294.979,
                4452.890,
                4616.482,
                4785.962,
                4961.548,
                5143.463,
                5331.939,
                5527.217,
                5729.545,
                5939.183,
                6156.396,
                6381.463,
                6614.671,
                6856.316,
                7106.708,
                7366.166,
                7635.020,
                7913.614,
                8202.302,
                8501.454,
                8811.450,
                9132.688,
                9465.574,
                9810.536,
                10168.013,
                10538.460,
                10922.351,
                11320.175,
                11732.438,
                12159.670,
                12602.412,
                13061.229,
                13536.710,
                14029.458,
                14540.103,
                15069.295,
                15617.710,
                16186.049,
                16775.035,
                17385.420,
            ]
        )
        f_c = np.array(
            [
                91.708,
                115.216,
                138.870,
                162.702,
                186.742,
                211.019,
                235.566,
                260.413,
                285.593,
                311.136,
                337.077,
                363.448,
                390.282,
                417.614,
                445.479,
                473.912,
                502.950,
                532.629,
                562.988,
                594.065,
                625.899,
                658.533,
                692.006,
                726.362,
                761.644,
                797.898,
                835.170,
                873.508,
                912.959,
                953.576,
                995.408,
                1038.511,
                1082.938,
                1128.746,
                1175.995,
                1224.744,
                1275.055,
                1326.992,
                1380.623,
                1436.014,
                1493.237,
                1552.366,
                1613.474,
                1676.641,
                1741.946,
                1809.474,
                1879.310,
                1951.543,
                2026.266,
                2103.573,
                2183.564,
                2266.340,
                2352.008,
                2440.675,
                2532.456,
                2627.468,
                2725.832,
                2827.672,
                2933.120,
                3042.309,
                3155.379,
                3272.475,
                3393.745,
                3519.344,
                3649.432,
                3784.176,
                3923.748,
                4068.324,
                4218.090,
                4373.237,
                4533.963,
                4700.473,
                4872.978,
                5051.700,
                5236.866,
                5428.712,
                5627.484,
                5833.434,
                6046.825,
                6267.931,
                6497.031,
                6734.420,
                6980.399,
                7235.284,
                7499.397,
                7773.077,
                8056.673,
                8350.547,
                8655.072,
                8970.639,
                9297.648,
                9636.520,
                9987.683,
                10351.586,
                10728.695,
                11119.490,
                11524.470,
                11944.149,
                12379.066,
                12829.775,
                13294.850,
                13780.887,
                14282.503,
                14802.338,
                15341.057,
                15899.345,
                16477.914,
                17077.504,
                17690.045,
            ]
        )
        f_u = np.array(
            [
                103.445,
                127.023,
                150.762,
                174.694,
                198.849,
                223.257,
                247.950,
                272.959,
                298.317,
                324.055,
                350.207,
                376.805,
                403.884,
                431.478,
                459.622,
                488.353,
                517.707,
                547.721,
                578.434,
                609.885,
                642.114,
                675.161,
                709.071,
                743.884,
                779.647,
                816.404,
                854.203,
                893.091,
                933.113,
                974.336,
                1016.797,
                1060.555,
                1105.666,
                1152.187,
                1200.178,
                1249.700,
                1300.816,
                1353.592,
                1408.094,
                1464.392,
                1522.559,
                1582.668,
                1644.795,
                1709.021,
                1775.427,
                1844.098,
                1915.121,
                1988.587,
                2064.590,
                2143.227,
                2224.597,
                2308.806,
                2395.959,
                2486.169,
                2579.551,
                2676.223,
                2776.309,
                2879.937,
                2987.238,
                3098.350,
                3213.415,
                3332.579,
                3455.993,
                3583.817,
                3716.212,
                3853.348,
                3995.399,
                4142.547,
                4294.979,
                4452.890,
                4643.482,
                4785.962,
                4961.548,
                5143.463,
                5331.939,
                5527.217,
                5729.545,
                5939.183,
                6156.396,
                6381.463,
                6614.671,
                6856.316,
                7106.708,
                7366.166,
                7635.020,
                7913.614,
                8202.302,
                8501.454,
                8811.450,
                9132.688,
                9465.574,
                9810.536,
                10168.013,
                10538.460,
                10922.351,
                11320.175,
                11732.438,
                12159.670,
                12602.412,
                13061.229,
                13536.710,
                14029.458,
                14540.103,
                15069.295,
                15617.710,
                16186.049,
                16775.035,
                17385.420,
                18000.000,
            ]
        )
        df = self.sample_rate / self.window_length
        filterbank = torch.zeros(self.window_length // 2 + 1, len(f_c))
        for k in range(filterbank.shape[0]):
            for i in range(filterbank.shape[1]):
                temp = (
                    np.amin([f_u[i], (k + 0.5) * df])
                    - np.amax([f_l[i], (k - 0.5) * df])
                ) / df
                filterbank[k, i] = np.amax([0, temp])
        return filterbank, f_c

    def get_noise(self):
        """Internal ear noise"""
        return torch.tensor(
            [10 ** (1.456 * ((f / 1000) ** -0.8) / 10) for f in self.f_c]
        ).unsqueeze(0)

    def get_alpha(self):
        """Get alpha for time spreading"""
        tau_100 = 0.030
        tau_min = 0.008
        Nadv = self.window_length / 2
        Fss = self.sample_rate / Nadv
        tau = tau_min + (np.divide(100, self.f_c)) * (tau_100 - tau_min)
        alpha = np.exp(np.divide(-1 / Fss, tau))
        return torch.tensor(alpha).unsqueeze(0)

    def slicing(self, sequence: np.array):
        """
        Input
            sequence: numpy array with a shape [batch, length]
        Output
            numpy array with a shape [batch, num_chunks, window_length]
        """
        num_chunks = (sequence.shape[-1] - self.window_length) // (
            self.window_length - self.overlap
        ) + 1
        chunks = np.stack(
            [
                np.lib.stride_tricks.sliding_window_view(s, self.window_length)[
                    :: self.window_length - self.overlap
                ]
                for s in sequence
            ]
        )  # batch, num_slice, window_length
        return chunks[:, :num_chunks, :]

    def windowing(self, stacked_wav):
        """Hann window"""
        return stacked_wav * self.weighted_hann

    def DFT(self, windowed_wav):
        """Discrete Fourier transform"""
        effective_bins = self.window_length // 2
        spec = torch.fft.fft(torch.tensor(windowed_wav), n=self.window_length)
        squared_spec = torch.cat(
            [
                spec[:, :1].real ** 2,
                spec[:, 1:effective_bins].real ** 2
                + spec[:, 1:effective_bins].imag ** 2,
                spec[:, effective_bins : effective_bins + 1].real ** 2,
            ],
            dim=-1,
        )
        return squared_spec

    def frequency_spreading(self, E, Bs):
        """frequency spreading"""
        e = 0.4
        dz = 0.25
        aL = 10 ** (-2.7 * dz)
        aUC = torch.tensor(10 ** ((-2.4 - 23 / self.f_c) * dz))
        aUCE = aUC * (E ** (0.2 * dz))
        gIL = torch.tensor(
            [(1 - aL ** (i + 1)) / (1 - aL) for i in range(len(self.f_c))]
        ).unsqueeze(0)
        gIU = torch.stack(
            [
                (1 - aUCE[:, i] ** (len(self.f_c) - i)) / (1 - aUCE[:, i])
                for i in range(len(self.f_c))
            ]
        ).T
        En = E / (gIL + gIU - 1)
        aUCEe = aUCE**e
        Ene = En**e

        # lower spreading
        Es = torch.zeros(Ene.shape)
        Es[:, -1] = Ene[:, -1]
        aLe = aL**e
        for i in range(len(self.f_c) - 2, -1, -1):
            Es[:, i] = aLe * Es[:, i + 1] + Ene[:, i]

        # upper spreading
        for i in range(len(self.f_c) - 1):
            r = Ene[:, i]
            a = aUCEe[:, i]
            for j in range(i + 1, len(self.f_c)):
                r = r * a
                Es[:, j] += r

        Es = (Es ** (1 / e)) / Bs
        return Es

    def time_spreading(self, Es):
        """time spreading"""
        b, t, f = Es.shape
        prev = torch.zeros(b, f)
        for i in range(t):
            prev = self.alpha * prev + (1 - self.alpha) * Es[:, i, :]
            Es[:, i, :] = torch.max(prev, Es[:, i, :])
        return Es

    def adbb(self, E_test, E_ref):
        """average distorted block"""
        c = [-0.198719, 0.0550197, -0.00102438, 5.05622e-6, 9.01033e-11]
        d1 = 5.95072
        d2 = 6.39468
        g = 1.71332
        bP = 4
        bM = 6

        E_test = 10 * np.log10(E_test.numpy())
        E_ref = 10 * np.log10(E_ref.numpy())

        # asymmetric excitation
        E_diff = E_ref - E_test
        L = np.zeros(E_test.shape)
        L[E_diff > 0] = 0.3 * E_ref[E_diff > 0] + 0.7 * E_test[E_diff > 0]
        L[E_diff <= 0] = E_test[E_diff <= 0]

        # effective detection step size (just noticeable difference)
        s = np.zeros(L.shape)
        s[L > 0] = (
            c[0]
            + c[1] * L[L > 0]
            + c[2] * (L[L > 0] ** 2)
            + c[3] * (L[L > 0] ** 3)
            + c[4] * (L[L > 0] ** 4)
            + d1 * (d2 / L[L > 0]) ** g
        )
        s[L <= 0] = 10**30

        # detection probability above threshold
        p = np.zeros(s.shape)
        p[E_diff > 0] = 1 - 0.5 ** ((E_diff[E_diff > 0] / s[E_diff > 0]) ** bP)
        p[E_diff <= 0] = 1 - 0.5 ** ((E_diff[E_diff <= 0] / s[E_diff <= 0]) ** bM)
        q = np.abs(np.floor(E_diff)) / s

        # average
        counter = p[p > 0.5].sum()
        if counter == 0:
            adbb = 0.0
        else:
            adbb = np.log10(q[p > 0.5].sum() / p[p > 0.5].sum())
        return adbb

    def get_representation(self, audio: AudioSignal):
        """
        Input
            audio: AudioSignal
        """
        # input audio
        wav = audio.audio_data.squeeze(1).cpu().detach().numpy()
        b = len(wav)

        # slicing
        stacked_wav = self.slicing(wav)
        stacked_wav = rearrange(stacked_wav, "b t f -> (b t) f")

        # windowing
        windowed_wav = self.windowing(stacked_wav)

        # discrete Fourier transform
        squared_spec = self.DFT(windowed_wav)

        # outer and middle ear modeling
        weighted_spec = (self.omem_weights**2) * squared_spec

        # grouping filterbank
        grouped_spec = torch.matmul(weighted_spec, self.filterbank)
        grouped_spec[grouped_spec < 1e-12] = 1e-12

        # internal ear noise
        noised_spec = grouped_spec + self.internal_noise

        # frequency spreading
        freq_spread_spec = self.frequency_spreading(noised_spec, self.bs)

        # rearrange
        noised_repr = rearrange(noised_spec, "(b t) f -> b t f", b=b)
        freq_spread_repr = rearrange(freq_spread_spec, "(b t) f -> b t f", b=b)

        # time spreading
        time_spread_repr = self.time_spreading(freq_spread_repr)

        return (
            20 * torch.log10(noised_repr),
            20 * torch.log10(freq_spread_repr),
            # 20 * torch.log10(time_spread_repr),
            time_spread_repr,
        )

    def forward(self, x: AudioSignal, y: AudioSignal):
        """Computes mel loss between an estimate and a reference
        signal.

        Parameters
        ----------
        x : AudioSignal
            Estimate signal
        y : AudioSignal
            Reference signal

        Returns
        -------
        torch.Tensor
            Mel loss.
        """
        # process
        x_repr_noised, x_repr_spread, E_test = self.get_representation(x)
        y_repr_noised, y_repr_spread, E_ref = self.get_representation(y)

        # adbb
        # adbb = self.adbb(E_test, E_ref)

        loss_noised = (1 - self.ratio) * self.loss_fn(x_repr_noised, y_repr_noised)
        loss_spread = self.ratio * self.loss_fn(x_repr_spread, y_repr_spread)

        return loss_noised + loss_spread