import math
import warnings
from typing import List, Optional, Union

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

from mayavoz.data.dataset import MayaDataset
from mayavoz.models.model import Mayamodel
from mayavoz.utils.io import Audio as audio
from mayavoz.utils.utils import merge_dict


class DemucsLSTM(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int,
        bidirectional: bool = True,
    ):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size, hidden_size, num_layers, bidirectional=bidirectional
        )
        dim = 2 if bidirectional else 1
        self.linear = nn.Linear(dim * hidden_size, hidden_size)

    def forward(self, x):

        output, (h, c) = self.lstm(x)
        output = self.linear(output)

        return output, (h, c)


class DemucsEncoder(nn.Module):
    def __init__(
        self,
        num_channels: int,
        hidden_size: int,
        kernel_size: int,
        stride: int = 1,
        glu: bool = False,
    ):
        super().__init__()
        activation = nn.GLU(1) if glu else nn.ReLU()
        multi_factor = 2 if glu else 1
        self.encoder = nn.Sequential(
            nn.Conv1d(num_channels, hidden_size, kernel_size, stride),
            nn.ReLU(),
            nn.Conv1d(hidden_size, hidden_size * multi_factor, 1, 1),
            activation,
        )

    def forward(self, waveform):

        return self.encoder(waveform)


class DemucsDecoder(nn.Module):
    def __init__(
        self,
        num_channels: int,
        hidden_size: int,
        kernel_size: int,
        stride: int = 1,
        glu: bool = False,
        layer: int = 0,
    ):
        super().__init__()
        activation = nn.GLU(1) if glu else nn.ReLU()
        multi_factor = 2 if glu else 1
        self.decoder = nn.Sequential(
            nn.Conv1d(hidden_size, hidden_size * multi_factor, 1, 1),
            activation,
            nn.ConvTranspose1d(hidden_size, num_channels, kernel_size, stride),
        )
        if layer > 0:
            self.decoder.add_module("4", nn.ReLU())

    def forward(
        self,
        waveform,
    ):

        out = self.decoder(waveform)
        return out


class Demucs(Mayamodel):
    """
    Demucs model from https://arxiv.org/pdf/1911.13254.pdf
    parameters:
        encoder_decoder: dict, optional
            keyword arguments passsed to encoder decoder block
        lstm : dict, optional
            keyword arguments passsed to LSTM block
        num_channels: int, defaults to 1
            number channels in input audio
        sampling_rate: int, defaults to 16KHz
            sampling rate of input audio
        lr : float, defaults to 1e-3
            learning rate used for training
        dataset: MayaDataset, optional
            MayaDataset object containing train/validation data for training
        duration : float, optional
            chunk duration in seconds
        loss : string or List of strings
            loss function to be used, available ("mse","mae","SI-SDR")
        metric : string or List of strings
            metric function to be used, available ("mse","mae","SI-SDR")

    """

    ED_DEFAULTS = {
        "initial_output_channels": 48,
        "kernel_size": 8,
        "stride": 4,
        "depth": 5,
        "glu": True,
        "growth_factor": 2,
    }
    LSTM_DEFAULTS = {
        "bidirectional": True,
        "num_layers": 2,
    }

    def __init__(
        self,
        encoder_decoder: Optional[dict] = None,
        lstm: Optional[dict] = None,
        num_channels: int = 1,
        resample: int = 4,
        sampling_rate=16000,
        normalize=True,
        lr: float = 1e-3,
        dataset: Optional[MayaDataset] = None,
        duration: Optional[float] = None,
        loss: Union[str, List] = "mse",
        metric: Union[str, List] = "mse",
        floor=1e-3,
    ):
        duration = (
            dataset.duration if isinstance(dataset, MayaDataset) else duration
        )
        if dataset is not None:
            if sampling_rate != dataset.sampling_rate:
                warnings.warn(
                    f"model sampling rate {sampling_rate} should match dataset sampling rate {dataset.sampling_rate}"
                )
                sampling_rate = dataset.sampling_rate
        super().__init__(
            num_channels=num_channels,
            sampling_rate=sampling_rate,
            lr=lr,
            dataset=dataset,
            duration=duration,
            loss=loss,
            metric=metric,
        )

        encoder_decoder = merge_dict(self.ED_DEFAULTS, encoder_decoder)
        lstm = merge_dict(self.LSTM_DEFAULTS, lstm)
        self.save_hyperparameters("encoder_decoder", "lstm", "resample")
        hidden = encoder_decoder["initial_output_channels"]
        self.normalize = normalize
        self.floor = floor
        self.encoder = nn.ModuleList()
        self.decoder = nn.ModuleList()

        for layer in range(encoder_decoder["depth"]):

            encoder_layer = DemucsEncoder(
                num_channels=num_channels,
                hidden_size=hidden,
                kernel_size=encoder_decoder["kernel_size"],
                stride=encoder_decoder["stride"],
                glu=encoder_decoder["glu"],
            )
            self.encoder.append(encoder_layer)

            decoder_layer = DemucsDecoder(
                num_channels=num_channels,
                hidden_size=hidden,
                kernel_size=encoder_decoder["kernel_size"],
                stride=encoder_decoder["stride"],
                glu=encoder_decoder["glu"],
                layer=layer,
            )
            self.decoder.insert(0, decoder_layer)

            num_channels = hidden
            hidden = self.ED_DEFAULTS["growth_factor"] * hidden

        self.de_lstm = DemucsLSTM(
            input_size=num_channels,
            hidden_size=num_channels,
            num_layers=lstm["num_layers"],
            bidirectional=lstm["bidirectional"],
        )

    def forward(self, waveform):

        if waveform.dim() == 2:
            waveform = waveform.unsqueeze(1)

        if waveform.size(1) != self.hparams.num_channels:
            raise ValueError(
                f"Number of input channels initialized is {self.hparams.num_channels} but got {waveform.size(1)} channels"
            )
        if self.normalize:
            waveform = waveform.mean(dim=1, keepdim=True)
            std = waveform.std(dim=-1, keepdim=True)
            waveform = waveform / (self.floor + std)
        else:
            std = 1
        length = waveform.shape[-1]
        x = F.pad(waveform, (0, self.get_padding_length(length) - length))
        if self.hparams.resample > 1:
            x = audio.resample_audio(
                audio=x,
                sr=self.hparams.sampling_rate,
                target_sr=int(
                    self.hparams.sampling_rate * self.hparams.resample
                ),
            )

        encoder_outputs = []
        for encoder in self.encoder:
            x = encoder(x)
            encoder_outputs.append(x)
        x = x.permute(0, 2, 1)
        x, _ = self.de_lstm(x)

        x = x.permute(0, 2, 1)
        for decoder in self.decoder:
            skip_connection = encoder_outputs.pop(-1)
            x = x + skip_connection[..., : x.shape[-1]]
            x = decoder(x)

        if self.hparams.resample > 1:
            x = audio.resample_audio(
                x,
                int(self.hparams.sampling_rate * self.hparams.resample),
                self.hparams.sampling_rate,
            )

        out = x[..., :length]
        return std * out

    def get_padding_length(self, input_length):

        input_length = math.ceil(input_length * self.hparams.resample)

        for layer in range(
            self.hparams.encoder_decoder["depth"]
        ):  # encoder operation
            input_length = (
                math.ceil(
                    (input_length - self.hparams.encoder_decoder["kernel_size"])
                    / self.hparams.encoder_decoder["stride"]
                )
                + 1
            )
            input_length = max(1, input_length)
        for layer in range(
            self.hparams.encoder_decoder["depth"]
        ):  # decoder operaration
            input_length = (input_length - 1) * self.hparams.encoder_decoder[
                "stride"
            ] + self.hparams.encoder_decoder["kernel_size"]
        input_length = math.ceil(input_length / self.hparams.resample)

        return int(input_length)
