import torch
import librosa
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from audiotools import AudioSignal
from audiotools import ml
from audiotools import STFTParams
from einops import rearrange
from torch.nn.utils import weight_norm


def WNConv1d(*args, **kwargs):
    act = kwargs.pop("act", True)
    conv = weight_norm(nn.Conv1d(*args, **kwargs))
    if not act:
        return conv
    return nn.Sequential(conv, nn.LeakyReLU(0.1))


def WNConv2d(*args, **kwargs):
    act = kwargs.pop("act", True)
    conv = weight_norm(nn.Conv2d(*args, **kwargs))
    if not act:
        return conv
    return nn.Sequential(conv, nn.LeakyReLU(0.1))


class MPD(nn.Module):
    def __init__(self, period):
        super().__init__()
        self.period = period
        self.convs = nn.ModuleList(
            [
                WNConv2d(2, 32, (5, 1), (3, 1), padding=(2, 0)),
                WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
                WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
                WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
                WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
            ]
        )
        self.conv_post = WNConv2d(
            1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False
        )

    def pad_to_period(self, x):
        t = x.shape[-1]
        x = F.pad(x, (0, self.period - t % self.period), mode="reflect")
        return x

    def forward(self, x):
        fmap = []

        x = self.pad_to_period(x)
        x = rearrange(x, "b c (l p) -> b c l p", p=self.period)

        for layer in self.convs:
            x = layer(x)
            fmap.append(x)

        x = self.conv_post(x)
        fmap.append(x)

        return fmap


class MSD(nn.Module):
    def __init__(self, rate: int = 1, sample_rate: int = 44100):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                WNConv1d(2, 16, 15, 1, padding=7),
                WNConv1d(16, 64, 41, 4, groups=4, padding=20),
                WNConv1d(64, 256, 41, 4, groups=16, padding=20),
                WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
                WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
                WNConv1d(1024, 1024, 5, 1, padding=2),
            ]
        )
        self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False)
        self.sample_rate = sample_rate
        self.rate = rate

    def forward(self, x):
        x = AudioSignal(x, self.sample_rate)
        x.resample(self.sample_rate // self.rate)
        x = x.audio_data

        fmap = []

        for l in self.convs:
            x = l(x)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)

        return fmap


BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]


class MRD(nn.Module):
    def __init__(
        self,
        window_length: int,
        hop_factor: float = 0.25,
        sample_rate: int = 44100,
        bands: list = BANDS,
    ):
        """Complex multi-band spectrogram discriminator.
        Parameters
        ----------
        window_length : int
            Window length of STFT.
        hop_factor : float, optional
            Hop factor of the STFT, defaults to ``0.25 * window_length``.
        sample_rate : int, optional
            Sampling rate of audio in Hz, by default 44100
        bands : list, optional
            Bands to run discriminator over.
        """
        super().__init__()

        self.window_length = window_length
        self.hop_factor = hop_factor
        self.sample_rate = sample_rate
        self.stft_params = STFTParams(
            window_length=window_length,
            hop_length=int(window_length * hop_factor),
            match_stride=True,
        )

        n_fft = window_length // 2 + 1
        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
        self.bands = bands

        ch = 32
        convs = lambda: nn.ModuleList(
            [
                WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
            ]
        )
        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
        self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False)

    def spectrogram(self, x):
        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
        x = torch.view_as_real(x.stft())
        x = rearrange(x, "b ch f t c -> (b ch) c t f", ch=2)

        # Split into bands
        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
        return x_bands

    def forward(self, x):
        x_bands = self.spectrogram(x)
        fmap = []

        x = []
        for band, stack in zip(x_bands, self.band_convs):
            for layer in stack:
                band = layer(band)
                fmap.append(band)
            x.append(band)

        x = torch.cat(x, dim=-1)
        x = self.conv_post(x)
        fmap.append(x)

        return fmap


class SubbandProjection(nn.Module):
    def __init__(self, bandwidth, out_dim):
        super(SubbandProjection, self).__init__()
        self.layer_norm = nn.LayerNorm(bandwidth)
        self.fc = nn.Linear(bandwidth, out_dim)

    def forward(self, x):
        x = rearrange(x, "b f t -> b t f")
        x = self.layer_norm(x)
        x = self.fc(x)
        x = rearrange(x, "b t f -> b f t")
        return x


class NeuralFilterbank(nn.Module):
    def __init__(
        self, n_fft: int, n_filterbank: int, sample_rate: int = 48000, out_dim: int = 64
    ):
        super(NeuralFilterbank, self).__init__()

        self.bandwidth_indices = self.get_bandwidth_indices(
            sample_rate, n_fft, n_filterbank
        )
        self.projection_modules = self.get_projection_layers(out_dim)

    def get_bandwidth_indices(self, sample_rate, n_fft, n_filterbank):
        mel_basis = librosa.filters.mel(
            sr=sample_rate, n_fft=n_fft, n_mels=n_filterbank
        )
        bandwidth_indices = [np.where(row > 0)[0] for row in mel_basis]
        return bandwidth_indices

    def get_projection_layers(self, out_dim):
        projection_modules = nn.ModuleList([])
        for indices in self.bandwidth_indices:
            projection_modules.append(SubbandProjection(len(indices), out_dim))
        return projection_modules

    def forward(self, spec):
        # spec: [batch, freq, time]
        emb = []
        for indices, layer in zip(self.bandwidth_indices, self.projection_modules):
            emb.append(layer(spec[:, indices, :]))
        return rearrange(torch.stack(emb), "f b c t -> b c f t")


class Res2DMaxPoolModule(nn.Module):
    def __init__(self, input_channels, output_channels, pooling=2):
        super(Res2DMaxPoolModule, self).__init__()
        self.conv_1 = nn.Conv2d(input_channels, output_channels, 3, padding=1)
        self.bn_1 = nn.BatchNorm2d(output_channels)
        self.conv_2 = nn.Conv2d(output_channels, output_channels, 3, padding=1)
        self.bn_2 = nn.BatchNorm2d(output_channels)
        self.relu = nn.ReLU()
        self.mp = nn.MaxPool2d(pooling)

        # residual
        self.diff = False
        if input_channels != output_channels:
            self.conv_3 = nn.Conv2d(input_channels, output_channels, 3, padding=1)
            self.bn_3 = nn.BatchNorm2d(output_channels)
            self.diff = True

    def forward(self, x):
        out = self.bn_2(self.conv_2(self.relu(self.bn_1(self.conv_1(x)))))
        if self.diff:
            x = self.bn_3(self.conv_3(x))
        out = x + out
        out = self.mp(self.relu(out))
        return out


class ResFrontEnd(nn.Module):
    """
    Evaluation of CNN based Music Tagging.
    Won et al., 2020

    Note that, different from the original work, we only stack 3 convolutional layers instead of 7.
    After the convolution layers, we flatten the time-frequency representation to be a vector.
    """

    def __init__(self, conv_ndim, attention_ndim, nfreq, nchannel):
        super(ResFrontEnd, self).__init__()
        self.input_bn = nn.BatchNorm2d(nchannel)
        self.layer1 = Res2DMaxPoolModule(nchannel, conv_ndim, pooling=(2, 2))
        self.layer2 = Res2DMaxPoolModule(conv_ndim, conv_ndim, pooling=(2, 2))
        self.layer3 = Res2DMaxPoolModule(conv_ndim, conv_ndim, pooling=(2, 2))
        fc_ndim = nfreq // 2 // 2 // 2 * conv_ndim
        self.fc = nn.Linear(fc_ndim, attention_ndim)

    def forward(self, spec):
        # batch normalization
        out = self.input_bn(spec)

        # CNN
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)

        # permute and channel control
        b, c, f, t = out.shape
        out = out.permute(0, 3, 1, 2)  # batch, time, conv_ndim, freq
        out = out.contiguous().view(b, t, -1)  # batch, time, fc_ndim
        out = self.fc(out)  # batch, time, attention_ndim
        return out


# Transformer modules
"""
    Referenced PyTorch implementation of Vision Transformer by Lucidrains.
    https://github.com/lucidrains/vit-pytorch.git
"""


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x


class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        inner_dim = dim_head * heads
        self.heads = heads
        self.dropout = dropout
        self.scale = dim_head**-0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def forward(self, x, mask=None):
        b, n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)

        # # flash attention
        # with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
        #     out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout, is_causal=False)

        dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale
        mask_value = -torch.finfo(dots.dtype).max

        if mask is not None:
            mask = F.pad(mask.flatten(1), (1, 0), value=True)
            assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions"
            mask = mask[:, None, :] * mask[:, :, None]
            dots.masked_fill_(~mask, mask_value)
            del mask

        attn = dots.softmax(dim=-1)

        out = torch.einsum("bhij,bhjd->bhid", attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.to_out(out)
        return out


class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        Residual(
                            PreNorm(
                                dim,
                                Attention(
                                    dim, heads=heads, dim_head=dim_head, dropout=dropout
                                ),
                            )
                        ),
                        Residual(
                            PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))
                        ),
                    ]
                )
            )

    def forward(self, x, mask=None):
        for attn, ff in self.layers:
            x = attn(x, mask=mask)
            x = ff(x)
        return x


class ConvTransformer(nn.Module):
    def __init__(
        self,
        conv_ndim=16,
        n_channels=32,
        n_bands=32,
        attention_ndim=256,
        attention_nheads=8,
        attention_nlayers=4,
        attention_max_len=512,
        dropout=0.1,
    ):
        super(ConvTransformer, self).__init__()

        # Input embedding
        self.frontend = ResFrontEnd(conv_ndim, attention_ndim, n_bands, n_channels)

        # Positional embedding
        self.pos_embedding = nn.Parameter(
            torch.randn(1, attention_max_len + 1, attention_ndim)
        )

        # transformer
        self.transformer = Transformer(
            attention_ndim,
            attention_nlayers,
            attention_nheads,
            attention_ndim // attention_nheads,
            attention_ndim * 4,
            dropout,
        )
        self.to_latent = nn.Identity()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor):
        """

        Args:
            x (torch.Tensor): (batch, channel, frequency, time)

        Returns:
            x (torch.Tensor): (batch, out_ndim)

        """

        # Input embedding
        x = self.frontend(x)

        # Positional embedding with a [CLS] token
        x += self.pos_embedding[:, : x.size(1)]
        x = self.dropout(x)

        # transformer
        x = self.transformer(x)

        return x


class BSDConv(nn.Module):
    def __init__(
        self,
        n_bands: int,
        window_length: int,
        sample_rate: int = 48000,
    ):
        """Band-split spectrogram discriminator.
        Parameters
        ----------
        n_bands : int
            Number of frequency bands.
        window_length : int
            Window length of STFT.
        sample_rate : int, optional
            Sampling rate of audio in Hz, by default 44100
        """
        super().__init__()

        self.window_length = window_length
        self.n_bands = n_bands
        self.sample_rate = sample_rate
        self.stft_params = STFTParams(
            window_length=window_length,
            hop_length=window_length // 4,
            match_stride=True,
        )

        ch = 32
        self.neural_filterbank = NeuralFilterbank(
            n_fft=window_length,
            n_filterbank=n_bands,
            sample_rate=sample_rate,
            out_dim=ch,
        )
        self.convs = nn.ModuleList(
            [
                WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
                WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
                WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False),
            ]
        )

    def bs_spectrogram(self, x):
        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
        x = x.stft().abs().clamp(1e-5).pow(2.0).log10()
        x = rearrange(x, "b c f t -> (b c) f t")

        return self.neural_filterbank(x)

    def forward(self, x):
        fmap = []
        x = self.bs_spectrogram(x)
        for layer in self.convs:
            x = layer(x)
            fmap.append(x)
        return fmap


class BSDTransformer(nn.Module):
    def __init__(
        self,
        n_bands: int,
        window_length: int,
        sample_rate: int = 48000,
        conv_ndim: int = 16,
        attention_ndim: int = 32,
        attention_nheads: int = 2,
        attention_nlayers: int = 1,
    ):
        """Band-split spectrogram discriminator.
        Parameters
        ----------
        n_bands : int
            Number of frequency bands.
        window_length : int
            Window length of STFT.
        sample_rate : int, optional
            Sampling rate of audio in Hz, by default 44100
        """
        super().__init__()

        self.window_length = window_length
        self.n_bands = n_bands
        self.sample_rate = sample_rate
        self.stft_params = STFTParams(
            window_length=window_length,
            hop_length=window_length // 4,
            match_stride=True,
        )

        ch = 32
        self.neural_filterbank = NeuralFilterbank(
            n_fft=window_length,
            n_filterbank=n_bands,
            sample_rate=sample_rate,
            out_dim=ch,
        )
        self.transformer = ConvTransformer(
            conv_ndim, ch, n_bands, attention_ndim, attention_nheads, attention_nlayers
        )

    def bs_spectrogram(self, x):
        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
        x = x.stft().abs().clamp(1e-5).pow(2.0).log10()
        x = rearrange(x, "b c f t -> (b c) f t")

        return self.neural_filterbank(x)

    def forward(self, x):
        fmap = []
        x = self.bs_spectrogram(x)
        x = self.transformer(x)
        fmap.append(x)
        return fmap


class Res2dModule(nn.Module):
    def __init__(self, idim, odim, stride=(2, 2)):
        super(Res2dModule, self).__init__()
        self.conv1 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
        self.bn1 = nn.BatchNorm2d(odim)
        self.conv2 = nn.Conv2d(odim, odim, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(odim)
        self.relu = nn.ReLU()

        # residual
        self.diff = False
        if (idim != odim) or (stride[0] > 1):
            self.conv3 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
            self.bn3 = nn.BatchNorm2d(odim)
            self.diff = True

    def forward(self, x):
        out = self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))
        if self.diff:
            x = self.bn3(self.conv3(x))
        out = x + out
        out = self.relu(out)
        return out


class BSD(nn.Module):
    def __init__(
        self,
        n_bands: int,
        window_length: int,
        sample_rate: int = 48000,
    ):
        """Band-split spectrogram discriminator.
        Parameters
        ----------
        n_bands : int
            Number of frequency bands.
        window_length : int
            Window length of STFT.
        sample_rate : int, optional
            Sampling rate of audio in Hz, by default 44100
        """
        super().__init__()

        self.window_length = window_length
        self.n_bands = n_bands
        self.sample_rate = sample_rate
        self.stft_params = STFTParams(
            window_length=window_length,
            hop_length=window_length // 4,
            match_stride=True,
        )

        ch = 32
        self.neural_filterbank = NeuralFilterbank(
            n_fft=window_length,
            n_filterbank=n_bands,
            sample_rate=sample_rate,
            out_dim=ch,
        )
        self.convs = nn.ModuleList(
            [
                Res2dModule(ch, ch, (1, 1)),
                Res2dModule(ch, ch, (2, 2)),
                Res2dModule(ch, ch, (2, 2)),
                Res2dModule(ch, ch, (2, 2)),
                Res2dModule(ch, ch, (2, 2)),
                Res2dModule(ch, 1, (1, 1)),
            ]
        )

    def bs_spectrogram(self, x):
        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
        x = x.stft().abs().clamp(1e-5).pow(2.0).log10()
        x = rearrange(x, "b c f t -> (b c) f t")

        return self.neural_filterbank(x)

    def forward(self, x):
        fmap = []
        x = self.bs_spectrogram(x)
        fmap.append(x)
        for layer in self.convs:
            x = layer(x)
            fmap.append(x)
        return fmap


class Discriminator(ml.BaseModel):
    def __init__(
        self,
        rates: list = [],
        periods: list = [2, 3, 5, 7, 11],
        fft_sizes: list = [2048, 1024, 512],
        sample_rate: int = 48000,
        bands: list = BANDS,
    ):
        """Discriminator that combines multiple discriminators.

        Parameters
        ----------
        rates : list, optional
            sampling rates (in Hz) to run MSD at, by default []
            If empty, MSD is not used.
        periods : list, optional
            periods (of samples) to run MPD at, by default [2, 3, 5, 7, 11]
        fft_sizes : list, optional
            Window sizes of the FFT to run MRD at, by default [2048, 1024, 512]
        sample_rate : int, optional
            Sampling rate of audio in Hz, by default 44100
        bands : list, optional
            Bands to run MRD at, by default `BANDS`
        """
        super().__init__()
        discs = []
        discs += [MPD(p) for p in periods]
        discs += [MSD(r, sample_rate=sample_rate) for r in rates]
        discs += [MRD(f, sample_rate=sample_rate, bands=bands) for f in fft_sizes]
        # n_bands = [64, 128, 128, 160, 128, 160]
        # n_ffts = [512, 1024, 2048, 2048, 4096, 4096]
        n_bands = [160]
        n_ffts = [2048]
        bsds = []
        for f, b in zip(n_ffts, n_bands):
            bsds.append(BSD(b, f))
            bsds.append(BSDConv(b, f))
        discs += bsds
        self.discriminators = nn.ModuleList(discs)

    def preprocess(self, y):
        # Remove DC offset
        y = y - y.mean(dim=-1, keepdims=True)
        # Peak normalize the volume of input audio
        y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
        return y

    def forward(self, x):
        x = self.preprocess(x)
        fmaps = [d(x) for d in self.discriminators]
        return fmaps


if __name__ == "__main__":
    disc = Discriminator()
    x = torch.zeros(1, 2, 48000)
    results = disc(x)
    for i, result in enumerate(results):
        print(f"disc{i}")
        for i, r in enumerate(result):
            print(r.shape, r.mean(), r.min(), r.max())
        print()