import torch
import torch.nn as nn
import torch.nn.utils as utils

def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1):
    padding = (kernel_size // 2) * dilation
    return nn.Sequential(
        utils.spectral_norm(
            nn.Conv1d(in_channels, out_channels,
                kernel_size=kernel_size,
                stride=stride,
                dilation=dilation,
                padding=padding
            )
        ),
        nn.BatchNorm1d(out_channels),
        nn.LeakyReLU(0.2, inplace=True)
    )

class SISUDiscriminator(nn.Module):
    def __init__(self):
        super(SISUDiscriminator, self).__init__()
        layers = 4
        self.model = nn.Sequential(
            discriminator_block(1, layers, kernel_size=7, stride=2, dilation=1),
            discriminator_block(layers, layers * 2, kernel_size=5, stride=2, dilation=1),
            discriminator_block(layers * 2, layers * 4, kernel_size=3, dilation=4),
            discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=8),
            discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=16),
            discriminator_block(layers * 2, layers, kernel_size=5, dilation=2),
            discriminator_block(layers, 1, kernel_size=3, stride=1)
        )
        self.global_avg_pool = nn.AdaptiveAvgPool1d(1)

    def forward(self, x):
        x = self.model(x)
        x = self.global_avg_pool(x)
        return x.view(-1, 1)