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


class NegativeSNRLoss(nn.Module):
    """
    Signal-to-Noise Ratio
    """
    def __init__(self, eps: float = 1e-8):
        super(NegativeSNRLoss, self).__init__()
        self.eps = eps

    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
        """
        Compute the SI-SNR loss between the estimated signal and the target signal.

        :param denoise: The estimated signal (batch_size, signal_length)
        :param clean: The target signal (batch_size, signal_length)
        :return: The SI-SNR loss (batch_size,)
        """
        if denoise.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")

        denoise = denoise - torch.mean(denoise, dim=-1, keepdim=True)
        clean = clean - torch.mean(clean, dim=-1, keepdim=True)

        noise = denoise - clean

        clean_power = torch.norm(clean, p=2, dim=-1) ** 2
        noise_power = torch.norm(noise, p=2, dim=-1) ** 2

        snr = 10 * torch.log10((clean_power + self.eps) / (noise_power + self.eps))

        return -snr.mean()


class NegativeSISNRLoss(nn.Module):
    """
    Scale-Invariant Source-to-Noise Ratio

    https://arxiv.org/abs/2206.07293
    """
    def __init__(self,
                 reduction: str = "mean",
                 eps: float = 1e-8,
                 ):
        super(NegativeSISNRLoss, self).__init__()
        self.reduction = reduction
        self.eps = eps

    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
        """
        Compute the SI-SNR loss between the estimated signal and the target signal.

        :param denoise: The estimated signal (batch_size, signal_length)
        :param clean: The target signal (batch_size, signal_length)
        :return: The SI-SNR loss (batch_size,)
        """
        if denoise.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")

        denoise = denoise - torch.mean(denoise, dim=-1, keepdim=True)
        clean = clean - torch.mean(clean, dim=-1, keepdim=True)

        s_target = torch.sum(denoise * clean, dim=-1, keepdim=True) * clean / (torch.norm(clean, p=2, dim=-1, keepdim=True) ** 2 + self.eps)

        e_noise = denoise - s_target

        batch_si_snr = 10 * torch.log10(torch.norm(s_target, p=2, dim=-1) ** 2 / (torch.norm(e_noise, p=2, dim=-1) ** 2 + self.eps) + self.eps)
        # si_snr shape: [batch_size,]

        if self.reduction == "mean":
            loss = torch.mean(batch_si_snr)
        elif self.reduction == "sum":
            loss = torch.sum(batch_si_snr)
        else:
            raise AssertionError
        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)
    target_signal = torch.zeros(batch_size, signal_length)

    si_snr_loss = NegativeSISNRLoss()

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

    return


if __name__ == "__main__":
    main()