from typing import Optional

import torch
import numpy as np
from librosa.filters import mel

from .stft import STFT


class MelSpectrogram(torch.nn.Module):
    def __init__(
        self,
        is_half: bool,
        n_mel_channels: int,
        sampling_rate: int,
        win_length: int,
        hop_length: int,
        n_fft: Optional[int] = None,
        mel_fmin: int = 0,
        mel_fmax: int = None,
        clamp: float = 1e-5,
        device=torch.device("cpu"),
    ):
        super().__init__()
        if n_fft is None:
            n_fft = win_length
        mel_basis = mel(
            sr=sampling_rate,
            n_fft=n_fft,
            n_mels=n_mel_channels,
            fmin=mel_fmin,
            fmax=mel_fmax,
            htk=True,
        )
        mel_basis = torch.from_numpy(mel_basis).float()
        self.register_buffer("mel_basis", mel_basis)
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.clamp = clamp
        self.is_half = is_half

        self.stft = STFT(
            filter_length=n_fft,
            hop_length=hop_length,
            win_length=win_length,
            window="hann",
            use_torch_stft="privateuseone" not in str(device),
        ).to(device)

    def forward(
        self,
        audio: torch.Tensor,
        keyshift=0,
        speed=1,
        center=True,
    ):
        factor = 2 ** (keyshift / 12)
        win_length_new = int(np.round(self.win_length * factor))
        magnitude = self.stft(audio, keyshift, speed, center)
        if keyshift != 0:
            size = self.n_fft // 2 + 1
            resize = magnitude.size(1)
            if resize < size:
                magnitude = torch.nn.functional.pad(magnitude, (0, 0, 0, size - resize))
            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
        mel_output = torch.matmul(self.mel_basis, magnitude)
        if self.is_half:
            mel_output = mel_output.half()
        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
        return log_mel_spec