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


def GroupNormWrapper(generated, *args, **kwargs):
    """
    Wrapper around the group normalization layer generated by an instrument embedding, and standard group normalization

    Arguments:
        generated {bool} -- True if you want to use the generated groupnorm
        *args -- Positional arguments passed to the __init__ function of the chosen module
        **kwargs -- Keyword arguments passed to the __init__ function of the chosen module

    Returns:
        nn.Module
    """
    if generated: return GroupNormGenerated(*args)
    else:         return GroupNormStatic(*args)


class GroupNormGenerated(nn.Module):
    """
    Group normalization layer with scale and bias factor created with a linear transformation of the instrument embedding
    """
    def __init__(self, E_1, E_2, num_groups, num_channels, eps=1e-8):
        """
        Arguments:
            E_1 {int} -- Dimension of the instrument embedding
            E_2 {int} -- Dimension of the instrument embedding bottleneck
            num_groups {int} -- Number of normalized groups
            num_channels {int} -- Number of channels

        Keyword Arguments:
            eps {int} -- Constant for numerical stability (default: {1e-8})
        """
        super(GroupNormGenerated, self).__init__()

        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps

        self.bottleneck = nn.Linear(E_1, E_2)
        self.affine = nn.Linear(E_2, num_channels + num_channels)

    def forward(self, instrument, x):
        """
        Arguments:
            instrument {torch.tensor} -- Instrument embedding of shape (4, E_1)
            x {torch.tensor} -- Input of the groupnorm of shape (B, 4, C, T)

        Returns:
            torch.tensor -- Output of the groupnorm of shape (B, 4, C, T)
        """
        batch_size = x.shape[0]

        instrument = self.bottleneck(instrument)  # shape: (4, E_2)
        affine = self.affine(instrument)  # shape: (4, 2*C)

        scale = affine[:, :self.num_channels].contiguous().view(-1)  # shape: (4*C)
        bias = affine[:, self.num_channels:].contiguous().view(-1)  # shape: (4*C)

        x = x.view(batch_size, 4*self.num_channels, -1)  # shape: (B, 4*C, T)
        x = F.group_norm(x, 4*self.num_groups, scale, bias, self.eps)  # shape: (B, 4*C, T)
        x = x.view(batch_size, 4, self.num_channels, -1)  # shape: (B, 4, C, T)

        return x  # shape: (B, 4, C, T)


class GroupNormStatic(nn.Module):
    """
    Group normalization layer with an independent scale and bias factor for each instrument
    """
    def __init__(self, _, __, num_groups, num_channels, eps=1e-8):
        """
        Arguments:
            num_groups {int} -- Number of normalized groups
            num_channels {int} -- Number of channels

        Keyword Arguments:
            eps {int} -- Constant for numerical stability (default: {1e-8})
        """
        super(GroupNormStatic, self).__init__()

        self.num_groups = num_groups
        self.num_channels = num_channels
        self.group_norm = nn.GroupNorm(4*num_groups, 4*num_channels, eps)

    def forward(self, _, x):
        """
        Arguments:
            _ {None} -- unused argument (for compatibility with GroupNormGenerated)
            x {torch.tensor} -- Input of the groupnorm of shape (B, 4, C, T)

        Returns:
            torch.tensor -- Output of the groupnorm of shape (B, 4, C, T)
        """
        batch_size = x.shape[0]

        x = x.view(batch_size, 4*self.num_channels, -1)  # shape: (B, 4*C, T)
        x = self.group_norm(x)  # shape: (B, 4*C, T)
        x = x.view(batch_size, 4, self.num_channels, -1)  # shape: (B, 4, C, T)

        return x  # shape: (B, 4, C, T)
