from turtle import forward
from mlassistant.core import Model
from torch import nn
from mlassistant.gan.models import GAN, BaseGenerator, BaseDiscriminator

from .entropy_minimization import EntropyMinimization


class ConvBlock(nn.Module):

    def __init__(self, num_classes, ndf) -> None:
        super().__init__()
        self.conv2d = nn.Conv2d(num_classes, ndf, kernel_size=4, stride=2, padding=1),
        self.activation = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        x = self.conv2d(x)
        return self.activation(x)

class Discriminator(BaseDiscriminator):  # the discriminator

    def __init__(self, num_classes, ndf=64) -> None:
        super().__init__()
        self.net = nn.Sequential(
            ConvBlock(num_classes, ndf), 
            ConvBlock(ndf, 2 * ndf), 
            ConvBlock(2 * ndf, 4 * ndf),
            ConvBlock(4 * ndf, 8 * ndf),
            ConvBlock(8 * ndf , 1),
            )

    def forward(self, x):
        return x


class Generator(BaseGenerator):  # segnet

    def __init__(self, conf):
        super().__init__(conf)
        self.net = EntropyMinimization()

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


class AdventGan(GAN):  # the gan

    def __init__(self, d_name: str):
        super().__init__(Generator(), {d_name: Discriminator()})
