import random

import torch
import torch.nn.functional as F


def stft(x, fft_size, hop_size, win_length, window):
    """Perform STFT and convert to magnitude spectrogram.
    Args:
        x (Tensor): Input signal tensor (B, T).
        fft_size (int): FFT size.
        hop_size (int): Hop size.
        win_length (int): Window length.
        window (str): Window function type.
    Returns:
        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
    """
    # TODO: work with complex here
    x_stft = torch.view_as_real(
        torch.stft(x, fft_size, hop_size, win_length, window, return_complex=True)
    )
    real = x_stft[..., 0]
    imag = x_stft[..., 1]
    # clamp is needed to avoid nan or inf
    return torch.sqrt(torch.clamp(real ** 2 + imag ** 2, min=1e-7)).transpose(2, 1)


def _static_stft(x, y, fft_size=1024, hop_size=120, win_length=600):
    window = torch.hann_window(win_length, device=x.device)
    x_mag = stft(x, fft_size, hop_size, win_length, window)
    y_mag = stft(y, fft_size, hop_size, win_length, window)
    del window
    return x_mag, y_mag


def _randomized_stft(x, y):
    # values are roughly taken from static MultiResolutionSTFTLoss
    fft_size = random.randint(512, 2048)
    win_length = random.randint(
        min(1000, max(200, int(round(fft_size/3)))),
        min(1400, max(300, int(round(fft_size/1.5))))
    )
    hop_size = random.randint(
        min(200, max(40, int(round(win_length/6)))),
        min(300, max(60, int(round(win_length/4))))
    )
    return _static_stft(x, y, fft_size=fft_size, hop_size=hop_size, win_length=win_length)


def stft_loss_rand(x, y, n_resolutions=3):
    """Randomized multi resolution STFT loss module."""
    sc_loss = 0.0
    mag_loss = 0.0
    for _ in range(n_resolutions):
        x_mag, y_mag = _randomized_stft(x, y)
        # spectral convergence loss
        sc_l = torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
        # magnitude loss
        mag_l = F.l1_loss(torch.log(y_mag), torch.log(x_mag))
        sc_loss += sc_l
        mag_loss += mag_l
        del x_mag, y_mag, sc_l, mag_l
    sc_loss /= n_resolutions
    mag_loss /= n_resolutions
    return sc_loss, mag_loss


def stft_loss(x, y):
    """Static multi resolution STFT loss module."""
    fft_sizes = [2048, 1024, 512]
    hop_sizes = [240, 120, 50]
    win_lengths = [1200, 600, 240]
    assert(len(fft_sizes) == len(hop_sizes) == len(win_lengths))
    n_resolutions = len(fft_sizes)
    sc_loss = 0.0
    mag_loss = 0.0
    for fft_size, hop_size, win_length in zip(fft_sizes, hop_sizes, win_lengths):
        x_mag, y_mag = _static_stft(
            x, y, fft_size=fft_size, hop_size=hop_size, win_length=win_length
        )
        # spectral convergence loss
        sc_l = torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
        # magnitude loss
        mag_l = F.l1_loss(torch.log(y_mag), torch.log(x_mag))
        sc_loss += sc_l
        mag_loss += mag_l
        del x_mag, y_mag, sc_l, mag_l
    sc_loss /= n_resolutions
    mag_loss /= n_resolutions
    return sc_loss, mag_loss