import torch
import auraloss
import resampy
import torchaudio
from pesq import pesq
import pyloudnorm as pyln


def crest_factor(x):
    """Compute the crest factor of waveform."""

    peak, _ = x.abs().max(dim=-1)
    rms = torch.sqrt((x ** 2).mean(dim=-1))

    return 20 * torch.log(peak / rms.clamp(1e-8))


def rms_energy(x):

    rms = torch.sqrt((x ** 2).mean(dim=-1))

    return 20 * torch.log(rms.clamp(1e-8))


def spectral_centroid(x):
    """Compute the crest factor of waveform.

    See: https://gist.github.com/endolith/359724

    """

    spectrum = torch.fft.rfft(x).abs()
    normalized_spectrum = spectrum / spectrum.sum()
    normalized_frequencies = torch.linspace(0, 1, spectrum.shape[-1])
    spectral_centroid = torch.sum(normalized_frequencies * normalized_spectrum)

    return spectral_centroid


def loudness(x, sample_rate):
    """Compute the loudness in dB LUFS of waveform."""
    meter = pyln.Meter(sample_rate)

    # add stereo dim if needed
    if x.shape[0] < 2:
        x = x.repeat(2, 1)

    return torch.tensor(meter.integrated_loudness(x.permute(1, 0).numpy()))


class MelSpectralDistance(torch.nn.Module):
    def __init__(self, sample_rate, length=65536):
        super().__init__()
        self.error = auraloss.freq.MelSTFTLoss(
            sample_rate,
            fft_size=length,
            hop_size=length,
            win_length=length,
            w_sc=0,
            w_log_mag=1,
            w_lin_mag=1,
            n_mels=128,
            scale_invariance=False,
        )

        # I think scale invariance may not work well,
        # since aspects of the phase may be considered?

    def forward(self, input, target):
        return self.error(input, target)


class PESQ(torch.nn.Module):
    def __init__(self, sample_rate):
        super().__init__()
        self.sample_rate = sample_rate

    def forward(self, input, target):
        if self.sample_rate != 16000:
            target = resampy.resample(
                target.view(-1).numpy(),
                self.sample_rate,
                16000,
            )
            input = resampy.resample(
                input.view(-1).numpy(),
                self.sample_rate,
                16000,
            )

        return pesq(
            16000,
            target,
            input,
            "wb",
        )


class CrestFactorError(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input, target):
        return torch.nn.functional.l1_loss(
            crest_factor(input),
            crest_factor(target),
        ).item()


class RMSEnergyError(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input, target):
        return torch.nn.functional.l1_loss(
            rms_energy(input),
            rms_energy(target),
        ).item()


class SpectralCentroidError(torch.nn.Module):
    def __init__(self, sample_rate, n_fft=2048, hop_length=512):
        super().__init__()

        self.spectral_centroid = torchaudio.transforms.SpectralCentroid(
            sample_rate,
            n_fft=n_fft,
            hop_length=hop_length,
        )

    def forward(self, input, target):
        return torch.nn.functional.l1_loss(
            self.spectral_centroid(input + 1e-16).mean(),
            self.spectral_centroid(target + 1e-16).mean(),
        ).item()


class LoudnessError(torch.nn.Module):
    def __init__(self, sample_rate: int, peak_normalize: bool = False):
        super().__init__()
        self.sample_rate = sample_rate
        self.peak_normalize = peak_normalize

    def forward(self, input, target):

        if self.peak_normalize:
            # peak normalize
            x = input / input.abs().max()
            y = target / target.abs().max()
        else:
            x = input
            y = target

        return torch.nn.functional.l1_loss(
            loudness(x.view(1, -1), self.sample_rate),
            loudness(y.view(1, -1), self.sample_rate),
        ).item()