#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://zhuanlan.zhihu.com/p/627039860

https://github.com/facebookresearch/denoiser/blob/main/denoiser/stft_loss.py
"""
from typing import List

import torch
import torch.nn as nn
from torch.nn import functional as F


class LSDLoss(nn.Module):
    """
    Log Spectral Distance

    Mean square error of power spectrum
    """
    def __init__(self,
                 n_fft: int = 512,
                 win_size: int = 512,
                 hop_size: int = 256,
                 center: bool = True,
                 eps: float = 1e-8,
                 reduction: str = "mean",
                 ):
        super(LSDLoss, self).__init__()
        self.n_fft = n_fft
        self.win_size = win_size
        self.hop_size = hop_size
        self.center = center
        self.eps = eps
        self.reduction = reduction

        if reduction not in ("sum", "mean"):
            raise AssertionError(f"param reduction must be sum or mean.")

    def forward(self, denoise_power: torch.Tensor, clean_power: torch.Tensor):
        """
        :param denoise_power: power spectrum of the estimated signal power spectrum  (batch_size, ...)
        :param clean_power: power spectrum of the target signal (batch_size, ...)
        :return:
        """
        denoise_power = denoise_power + self.eps
        clean_power = clean_power + self.eps

        log_denoise_power = torch.log10(denoise_power)
        log_clean_power = torch.log10(clean_power)

        # mean_square_error shape: [b, f]
        mean_square_error = torch.mean(torch.square(log_denoise_power - log_clean_power), dim=-1)

        if self.reduction == "mean":
            lsd_loss = torch.mean(mean_square_error)
        elif self.reduction == "sum":
            lsd_loss = torch.sum(mean_square_error)
        else:
            raise AssertionError
        return lsd_loss


class ComplexSpectralLoss(nn.Module):
    def __init__(self,
                 n_fft: int = 512,
                 win_size: int = 512,
                 hop_size: int = 256,
                 center: bool = True,
                 eps: float = 1e-8,
                 reduction: str = "mean",
                 factor_mag: float = 0.5,
                 factor_pha: float = 0.3,
                 factor_gra: float = 0.2,
                 ):
        super().__init__()
        self.n_fft = n_fft
        self.win_size = win_size
        self.hop_size = hop_size
        self.center = center
        self.eps = eps
        self.reduction = reduction

        self.factor_mag = factor_mag
        self.factor_pha = factor_pha
        self.factor_gra = factor_gra

        if reduction not in ("sum", "mean"):
            raise AssertionError(f"param reduction must be sum or mean.")

        self.window = nn.Parameter(torch.hann_window(win_size), requires_grad=False)

    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
        """
        :param denoise: The estimated signal (batch_size, signal_length)
        :param clean: The target signal (batch_size, signal_length)
        :return:
        """
        if denoise.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")

        # denoise_stft, clean_stft shape: [b, f, t]
        denoise_stft = torch.stft(
            denoise,
            n_fft=self.n_fft,
            win_length=self.win_size,
            hop_length=self.hop_size,
            window=self.window,
            center=self.center,
            pad_mode="reflect",
            normalized=False,
            return_complex=True
        )
        clean_stft = torch.stft(
            clean,
            n_fft=self.n_fft,
            win_length=self.win_size,
            hop_length=self.hop_size,
            window=self.window,
            center=self.center,
            pad_mode="reflect",
            normalized=False,
            return_complex=True
        )

        # complex_diff shape: [b, f, t], dtype: torch.complex64
        complex_diff = denoise_stft - clean_stft

        # magnitude_diff, phase_diff shape: [b, f, t], dtype: torch.float32
        magnitude_diff = torch.abs(complex_diff)
        phase_diff = torch.angle(complex_diff)

        # magnitude_loss, phase_loss shape: [b,]
        magnitude_loss = torch.norm(magnitude_diff, p=2, dim=(-1, -2))
        phase_loss = torch.norm(phase_diff, p=1, dim=(-1, -2))

        # phase_grad shape: [b, f, t-1], dtype: torch.float32
        phase_grad = torch.diff(torch.angle(denoise_stft), dim=-1)
        grad_loss = torch.mean(torch.abs(phase_grad), dim=(-1, -2))

        # loss, grad_loss shape: [b,]
        batch_loss = self.factor_mag * magnitude_loss + self.factor_pha * phase_loss + self.factor_gra * grad_loss
        # print(f"magnitude_loss: {magnitude_loss}")
        # print(f"phase_loss: {phase_loss}")
        # print(f"grad_loss: {grad_loss}")

        if self.reduction == "mean":
            loss = torch.mean(batch_loss)
        elif self.reduction == "sum":
            loss = torch.sum(batch_loss)
        else:
            raise AssertionError
        return loss


class SpectralConvergenceLoss(torch.nn.Module):
    """Spectral convergence loss module."""

    def __init__(self,
                 reduction: str = "mean",
                 eps: float = 1e-8,
                 ):
        super(SpectralConvergenceLoss, self).__init__()
        self.reduction = reduction
        self.eps = eps

        if reduction not in ("sum", "mean"):
            raise AssertionError(f"param reduction must be sum or mean.")

    def forward(self,
                denoise_magnitude: torch.Tensor,
                clean_magnitude: torch.Tensor,
                ):
        """
        :param denoise_magnitude: Tensor, shape: [batch_size, time_steps, freq_bins]
        :param clean_magnitude: Tensor, shape: [batch_size, time_steps, freq_bins]
        :return:
        """
        error_norm = torch.norm(denoise_magnitude - clean_magnitude, p="fro", dim=(-1, -2))
        truth_norm = torch.norm(clean_magnitude, p="fro", dim=(-1, -2))

        batch_loss = error_norm / (truth_norm + self.eps)

        if self.reduction == "mean":
            loss = torch.mean(batch_loss)
        elif self.reduction == "sum":
            loss = torch.sum(batch_loss)
        else:
            raise AssertionError

        return loss


class LogSTFTMagnitudeLoss(torch.nn.Module):
    """Log STFT magnitude loss module."""

    def __init__(self,
                 reduction: str = "mean",
                 eps: float = 1e-8,
                 ):
        super(LogSTFTMagnitudeLoss, self).__init__()
        self.reduction = reduction
        self.eps = eps

        if reduction not in ("sum", "mean"):
            raise AssertionError(f"param reduction must be sum or mean.")

    def forward(self,
                denoise_magnitude: torch.Tensor,
                clean_magnitude: torch.Tensor,
                ):
        """
        :param denoise_magnitude: Tensor, shape: [batch_size, time_steps, freq_bins]
        :param clean_magnitude: Tensor, shape: [batch_size, time_steps, freq_bins]
        :return:
        """

        loss = F.l1_loss(torch.log(denoise_magnitude + self.eps), torch.log(clean_magnitude + self.eps))

        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
            raise AssertionError("SpectralConvergenceLoss, nan or inf in loss")

        return loss


class STFTLoss(torch.nn.Module):
    """STFT loss module."""

    def __init__(self,
                 n_fft: int = 1024,
                 win_size: int = 600,
                 hop_size: int = 120,
                 center: bool = True,
                 reduction: str = "mean",
                 ):
        super(STFTLoss, self).__init__()
        self.n_fft = n_fft
        self.win_size = win_size
        self.hop_size = hop_size
        self.center = center
        self.reduction = reduction

        self.window = nn.Parameter(torch.hann_window(win_size), requires_grad=False)

        self.spectral_convergence_loss = SpectralConvergenceLoss(reduction=reduction)
        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss(reduction=reduction)

    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
        """
        :param denoise:
        :param clean:
        :return:
        """
        if denoise.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")

        # denoise_stft, clean_stft shape: [b, f, t]
        denoise_stft = torch.stft(
            denoise,
            n_fft=self.n_fft,
            win_length=self.win_size,
            hop_length=self.hop_size,
            window=self.window,
            center=self.center,
            pad_mode="reflect",
            normalized=False,
            return_complex=True
        )
        clean_stft = torch.stft(
            clean,
            n_fft=self.n_fft,
            win_length=self.win_size,
            hop_length=self.hop_size,
            window=self.window,
            center=self.center,
            pad_mode="reflect",
            normalized=False,
            return_complex=True
        )

        denoise_magnitude = torch.abs(denoise_stft)
        clean_magnitude = torch.abs(clean_stft)

        sc_loss = self.spectral_convergence_loss.forward(denoise_magnitude, clean_magnitude)
        mag_loss = self.log_stft_magnitude_loss.forward(denoise_magnitude, clean_magnitude)

        return sc_loss, mag_loss


class MultiResolutionSTFTLoss(torch.nn.Module):
    """Multi resolution STFT loss module."""

    def __init__(self,
                 fft_size_list: List[int] = None,
                 win_size_list: List[int] = None,
                 hop_size_list: List[int] = None,
                 factor_sc=0.1,
                 factor_mag=0.1,
                 reduction: str = "mean",
                 ):
        super(MultiResolutionSTFTLoss, self).__init__()
        fft_size_list = fft_size_list or [512, 1024, 2048]
        win_size_list = win_size_list or [240, 600, 1200]
        hop_size_list = hop_size_list or [50, 120, 240]

        if not len(fft_size_list) == len(win_size_list) == len(hop_size_list):
            raise AssertionError

        loss_fn_list = nn.ModuleList([])
        for n_fft, win_size, hop_size in zip(fft_size_list, win_size_list, hop_size_list):
            loss_fn_list.append(
                STFTLoss(
                    n_fft=n_fft,
                    win_size=win_size,
                    hop_size=hop_size,
                    reduction=reduction,
                )
            )

        self.loss_fn_list = loss_fn_list
        self.factor_sc = factor_sc
        self.factor_mag = factor_mag

    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
        """
        :param denoise:
        :param clean:
        :return:
        """
        if denoise.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")

        sc_loss = 0.0
        mag_loss = 0.0
        for loss_fn in self.loss_fn_list:
            sc_l, mag_l = loss_fn.forward(denoise, clean)
            sc_loss += sc_l
            mag_loss += mag_l
        sc_loss = sc_loss / len(self.loss_fn_list)
        mag_loss = mag_loss / len(self.loss_fn_list)

        sc_loss = self.factor_sc * sc_loss
        mag_loss = self.factor_mag * mag_loss

        loss = sc_loss + mag_loss
        return loss


def main():
    batch_size = 2
    signal_length = 16000
    estimated_signal = torch.randn(batch_size, signal_length)
    target_signal = torch.randn(batch_size, signal_length)

    # loss_fn = LSDLoss()
    # loss_fn = ComplexSpectralLoss()
    loss_fn = MultiResolutionSTFTLoss()

    loss = loss_fn.forward(estimated_signal, target_signal)
    print(f"loss: {loss.item()}")

    return


if __name__ == "__main__":
    main()