import numpy as np
from torch import nn


class GBlock(nn.Module):
    def __init__(self, in_channels, cond_channels, downsample_factor):
        super().__init__()

        self.in_channels = in_channels
        self.cond_channels = cond_channels
        self.downsample_factor = downsample_factor

        self.start = nn.Sequential(
            nn.AvgPool1d(downsample_factor, stride=downsample_factor),
            nn.ReLU(),
            nn.Conv1d(in_channels, in_channels * 2, kernel_size=3, padding=1),
        )
        self.lc_conv1d = nn.Conv1d(cond_channels, in_channels * 2, kernel_size=1)
        self.end = nn.Sequential(
            nn.ReLU(), nn.Conv1d(in_channels * 2, in_channels * 2, kernel_size=3, dilation=2, padding=2)
        )
        self.residual = nn.Sequential(
            nn.Conv1d(in_channels, in_channels * 2, kernel_size=1),
            nn.AvgPool1d(downsample_factor, stride=downsample_factor),
        )

    def forward(self, inputs, conditions):
        outputs = self.start(inputs) + self.lc_conv1d(conditions)
        outputs = self.end(outputs)
        residual_outputs = self.residual(inputs)
        outputs = outputs + residual_outputs

        return outputs


class DBlock(nn.Module):
    def __init__(self, in_channels, out_channels, downsample_factor):
        super().__init__()

        self.in_channels = in_channels
        self.downsample_factor = downsample_factor
        self.out_channels = out_channels

        self.donwsample_layer = nn.AvgPool1d(downsample_factor, stride=downsample_factor)
        self.layers = nn.Sequential(
            nn.ReLU(),
            nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(out_channels, out_channels, kernel_size=3, dilation=2, padding=2),
        )
        self.residual = nn.Sequential(
            nn.Conv1d(in_channels, out_channels, kernel_size=1),
        )

    def forward(self, inputs):
        if self.downsample_factor > 1:
            outputs = self.layers(self.donwsample_layer(inputs)) + self.donwsample_layer(self.residual(inputs))
        else:
            outputs = self.layers(inputs) + self.residual(inputs)
        return outputs


class ConditionalDiscriminator(nn.Module):
    def __init__(self, in_channels, cond_channels, downsample_factors=(2, 2, 2), out_channels=(128, 256)):
        super().__init__()

        assert len(downsample_factors) == len(out_channels) + 1

        self.in_channels = in_channels
        self.cond_channels = cond_channels
        self.downsample_factors = downsample_factors
        self.out_channels = out_channels

        self.pre_cond_layers = nn.ModuleList()
        self.post_cond_layers = nn.ModuleList()

        # layers before condition features
        self.pre_cond_layers += [DBlock(in_channels, 64, 1)]
        in_channels = 64
        for i, channel in enumerate(out_channels):
            self.pre_cond_layers.append(DBlock(in_channels, channel, downsample_factors[i]))
            in_channels = channel

        # condition block
        self.cond_block = GBlock(in_channels, cond_channels, downsample_factors[-1])

        # layers after condition block
        self.post_cond_layers += [
            DBlock(in_channels * 2, in_channels * 2, 1),
            DBlock(in_channels * 2, in_channels * 2, 1),
            nn.AdaptiveAvgPool1d(1),
            nn.Conv1d(in_channels * 2, 1, kernel_size=1),
        ]

    def forward(self, inputs, conditions):
        batch_size = inputs.size()[0]
        outputs = inputs.view(batch_size, self.in_channels, -1)
        for layer in self.pre_cond_layers:
            outputs = layer(outputs)
        outputs = self.cond_block(outputs, conditions)
        for layer in self.post_cond_layers:
            outputs = layer(outputs)

        return outputs


class UnconditionalDiscriminator(nn.Module):
    def __init__(self, in_channels, base_channels=64, downsample_factors=(8, 4), out_channels=(128, 256)):
        super().__init__()

        self.downsample_factors = downsample_factors
        self.in_channels = in_channels
        self.downsample_factors = downsample_factors
        self.out_channels = out_channels

        self.layers = nn.ModuleList()
        self.layers += [DBlock(self.in_channels, base_channels, 1)]
        in_channels = base_channels
        for i, factor in enumerate(downsample_factors):
            self.layers.append(DBlock(in_channels, out_channels[i], factor))
            in_channels *= 2
        self.layers += [
            DBlock(in_channels, in_channels, 1),
            DBlock(in_channels, in_channels, 1),
            nn.AdaptiveAvgPool1d(1),
            nn.Conv1d(in_channels, 1, kernel_size=1),
        ]

    def forward(self, inputs):
        batch_size = inputs.size()[0]
        outputs = inputs.view(batch_size, self.in_channels, -1)
        for layer in self.layers:
            outputs = layer(outputs)
        return outputs


class RandomWindowDiscriminator(nn.Module):
    """Random Window Discriminator as described in
    http://arxiv.org/abs/1909.11646"""

    def __init__(
        self,
        cond_channels,
        hop_length,
        uncond_disc_donwsample_factors=(8, 4),
        cond_disc_downsample_factors=((8, 4, 2, 2, 2), (8, 4, 2, 2), (8, 4, 2), (8, 4), (4, 2, 2)),
        cond_disc_out_channels=((128, 128, 256, 256), (128, 256, 256), (128, 256), (256,), (128, 256)),
        window_sizes=(512, 1024, 2048, 4096, 8192),
    ):
        super().__init__()
        self.cond_channels = cond_channels
        self.window_sizes = window_sizes
        self.hop_length = hop_length
        self.base_window_size = self.hop_length * 2
        self.ks = [ws // self.base_window_size for ws in window_sizes]

        # check arguments
        assert len(cond_disc_downsample_factors) == len(cond_disc_out_channels) == len(window_sizes)
        for ws in window_sizes:
            assert ws % hop_length == 0

        for idx, cf in enumerate(cond_disc_downsample_factors):
            assert np.prod(cf) == hop_length // self.ks[idx]

        # define layers
        self.unconditional_discriminators = nn.ModuleList([])
        for k in self.ks:
            layer = UnconditionalDiscriminator(
                in_channels=k, base_channels=64, downsample_factors=uncond_disc_donwsample_factors
            )
            self.unconditional_discriminators.append(layer)

        self.conditional_discriminators = nn.ModuleList([])
        for idx, k in enumerate(self.ks):
            layer = ConditionalDiscriminator(
                in_channels=k,
                cond_channels=cond_channels,
                downsample_factors=cond_disc_downsample_factors[idx],
                out_channels=cond_disc_out_channels[idx],
            )
            self.conditional_discriminators.append(layer)

    def forward(self, x, c):
        scores = []
        feats = []
        # unconditional pass
        for window_size, layer in zip(self.window_sizes, self.unconditional_discriminators):
            index = np.random.randint(x.shape[-1] - window_size)

            score = layer(x[:, :, index : index + window_size])
            scores.append(score)

        # conditional pass
        for window_size, layer in zip(self.window_sizes, self.conditional_discriminators):
            frame_size = window_size // self.hop_length
            lc_index = np.random.randint(c.shape[-1] - frame_size)
            sample_index = lc_index * self.hop_length
            x_sub = x[:, :, sample_index : (lc_index + frame_size) * self.hop_length]
            c_sub = c[:, :, lc_index : lc_index + frame_size]

            score = layer(x_sub, c_sub)
            scores.append(score)
        return scores, feats