import math
import torch
import librosa

# based on https://github.com/neuralaudio/hear-baseline/blob/main/hearbaseline/naive.py


class RandomMelProjection(torch.nn.Module):
    def __init__(
        self,
        sample_rate,
        embed_dim=4096,
        n_mels=128,
        n_fft=4096,
        hop_size=1024,
        seed=0,
        epsilon=1e-4,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.embed_dim = embed_dim
        self.n_mels = n_mels
        self.n_fft = n_fft
        self.hop_size = hop_size
        self.seed = seed
        self.epsilon = epsilon

        # Set random seed
        torch.random.manual_seed(self.seed)

        # Create a Hann window buffer to apply to frames prior to FFT.
        self.register_buffer("window", torch.hann_window(self.n_fft))

        # Create a mel filter buffer.
        mel_scale = torch.tensor(
            librosa.filters.mel(
                self.sample_rate,
                n_fft=self.n_fft,
                n_mels=self.n_mels,
            )
        )
        self.register_buffer("mel_scale", mel_scale)

        # Projection matrices.
        normalization = math.sqrt(self.n_mels)
        self.projection = torch.nn.Parameter(
            torch.rand(self.n_mels, self.embed_dim) / normalization,
            requires_grad=False,
        )

    def forward(self, x):
        bs, chs, samp = x.size()

        x = torch.stft(
            x.view(bs, -1),
            self.n_fft,
            self.hop_size,
            window=self.window,
            return_complex=True,
        )
        x = x.unsqueeze(1).permute(0, 1, 3, 2)

        # Apply the mel-scale filter to the power spectrum.
        x = torch.matmul(x.abs(), self.mel_scale.transpose(0, 1))

        # power scale
        x = torch.pow(x + self.epsilon, 0.3)

        # apply random projection
        e = x.matmul(self.projection)

        # take mean across temporal dim
        e = e.mean(dim=2).view(bs, -1)

        return e

    def compute_frame_embedding(self, x):
        # Compute the real-valued Fourier transform on windowed input signal.
        x = torch.fft.rfft(x * self.window)

        # Convert to a power spectrum.
        x = torch.abs(x) ** 2.0

        # Apply the mel-scale filter to the power spectrum.
        x = torch.matmul(x, self.mel_scale.transpose(0, 1))

        # Convert to a log mel spectrum.
        x = torch.log(x + self.epsilon)

        # Apply projection to get a 4096 dimension embedding
        embedding = x.matmul(self.projection)

        return embedding