# -*- coding: utf-8 -*-

# Copyright 2019 Tomoki Hayashi
#  MIT License (https://opensource.org/licenses/MIT)


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

from ..hparams import HParams


def _make_stft_cfg(hop_length, win_length=None):
    if win_length is None:
        win_length = 4 * hop_length
    n_fft = 2 ** (win_length - 1).bit_length()
    return dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)


def get_stft_cfgs(hp: HParams):
    assert hp.wav_rate == 44100, f"wav_rate must be 44100, got {hp.wav_rate}"
    return [_make_stft_cfg(h) for h in (100, 256, 512)]


def stft(x, n_fft, hop_length, win_length, window):
    dtype = x.dtype
    x = torch.stft(
        x.float(), n_fft, hop_length, win_length, window, return_complex=True
    )
    x = x.abs().to(dtype)
    x = x.transpose(2, 1)  # (b f t) -> (b t f)
    return x


class SpectralConvergengeLoss(nn.Module):
    def forward(self, x_mag, y_mag):
        """Calculate forward propagation.
        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            Tensor: Spectral convergence loss value.
        """
        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")


class LogSTFTMagnitudeLoss(nn.Module):
    def forward(self, x_mag, y_mag):
        """Calculate forward propagation.
        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            Tensor: Log STFT magnitude loss value.
        """
        return F.l1_loss(torch.log1p(x_mag), torch.log1p(y_mag))


class STFTLoss(nn.Module):
    def __init__(self, hp, stft_cfg: dict, window="hann_window"):
        super().__init__()
        self.hp = hp
        self.stft_cfg = stft_cfg
        self.spectral_convergenge_loss = SpectralConvergengeLoss()
        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss()
        self.register_buffer(
            "window", getattr(torch, window)(stft_cfg["win_length"]), persistent=False
        )

    def forward(self, x, y):
        """Calculate forward propagation.
        Args:
            x (Tensor): Predicted signal (B, T).
            y (Tensor): Groundtruth signal (B, T).
        Returns:
            Tensor: Spectral convergence loss value.
            Tensor: Log STFT magnitude loss value.
        """
        stft_cfg = dict(self.stft_cfg)
        x_mag = stft(x, **stft_cfg, window=self.window)  # (b t) -> (b t f)
        y_mag = stft(y, **stft_cfg, window=self.window)
        sc_loss = self.spectral_convergenge_loss(x_mag, y_mag)
        mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)
        return dict(sc=sc_loss, mag=mag_loss)


class MRSTFTLoss(nn.Module):
    def __init__(self, hp: HParams, window="hann_window"):
        """Initialize Multi resolution STFT loss module.
        Args:
            resolutions (list): List of (FFT size, hop size, window length).
            window (str): Window function type.
        """
        super().__init__()
        stft_cfgs = get_stft_cfgs(hp)
        self.stft_losses = nn.ModuleList()
        self.hp = hp
        for c in stft_cfgs:
            self.stft_losses += [STFTLoss(hp, c, window=window)]

    def forward(self, x, y):
        """Calculate forward propagation.
        Args:
            x (Tensor): Predicted signal (b t).
            y (Tensor): Groundtruth signal (b t).
        Returns:
            Tensor: Multi resolution spectral convergence loss value.
            Tensor: Multi resolution log STFT magnitude loss value.
        """
        assert (
            x.dim() == 2 and y.dim() == 2
        ), f"(b t) is expected, but got {x.shape} and {y.shape}."

        dtype = x.dtype

        x = x.float()
        y = y.float()

        # Align length
        x = x[..., : y.shape[-1]]
        y = y[..., : x.shape[-1]]

        losses = {}

        for f in self.stft_losses:
            d = f(x, y)
            for k, v in d.items():
                losses.setdefault(k, []).append(v)

        for k, v in losses.items():
            losses[k] = torch.stack(v, dim=0).mean().to(dtype)

        return losses