model_implementation

import torch
import torch.nn as nn

class Generator(nn.Module):
    """
    DCGAN Generator Network
    
    The generator maps a latent space vector (z) to an image.
    It uses a series of fractionally-strided convolutions (also called transposed convolutions)
    to upsample the input vector to the desired image size.
    
    Architecture follows the guidelines from the DCGAN paper:
    - No fully connected layers except the initial projection
    - Uses BatchNorm in all layers except the output layer
    - ReLU activations in all layers except the output (which uses Tanh)
    """
    def __init__(self, latent_dim=100, channels=3, img_size=64):
        super(Generator, self).__init__()
        self.latent_dim = latent_dim
        self.channels = channels
        self.img_size = img_size
        
        # Calculate initial feature map size
        self.init_size = img_size // 16  # For 64x64 images, this is 4
        self.projection = nn.Sequential(
            # Input is latent vector z
            nn.Linear(latent_dim, 1024 * self.init_size * self.init_size),
            nn.BatchNorm1d(1024 * self.init_size * self.init_size),
            nn.ReLU(True)
        )
        
        # Reshape to multiple feature maps
        self.conv_blocks = nn.Sequential(
            # Input: (batch, 1024, 4, 4)
            nn.ConvTranspose2d(1024, 512, 4, stride=2, padding=1),  # (batch, 512, 8, 8)
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1),   # (batch, 256, 16, 16)
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),   # (batch, 128, 32, 32)
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(128, channels, 4, stride=2, padding=1),  # (batch, channels, 64, 64)
            nn.Tanh()  # Output values between -1 and 1
        )
    
    def forward(self, z):
        """
        Forward pass of the generator
        
        Args:
            z (torch.Tensor): Batch of latent vectors with shape (batch_size, latent_dim)
            
        Returns:
            torch.Tensor: Generated images with shape (batch_size, channels, img_size, img_size)
        """
        # Project and reshape
        out = self.projection(z)
        out = out.view(out.size(0), 1024, self.init_size, self.init_size)
        
        # Apply convolutional blocks
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Module):
    """
    DCGAN Discriminator Network
    
    The discriminator classifies whether an image is real or generated (fake).
    It uses a series of strided convolutions to downsample the input image
    and classify it as real or fake.
    
    Architecture follows the guidelines from the DCGAN paper:
    - No pooling layers, only strided convolutions
    - No fully connected layers except for the final classification layer
    - BatchNorm in all layers except the first input layer and the output layer
    - LeakyReLU activations in all layers
    """
    def __init__(self, channels=3, img_size=64):
        super(Discriminator, self).__init__()
        self.channels = channels
        self.img_size = img_size
        
        self.model = nn.Sequential(
            # Input: (batch, channels, img_size, img_size)
            nn.Conv2d(channels, 64, 4, stride=2, padding=1),  # (batch, 64, 32, 32)
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(64, 128, 4, stride=2, padding=1),  # (batch, 128, 16, 16)
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(128, 256, 4, stride=2, padding=1),  # (batch, 256, 8, 8)
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(256, 512, 4, stride=2, padding=1),  # (batch, 512, 4, 4)
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(512, 1, 4, stride=1, padding=0)  # (batch, 1, 1, 1)
        )
    
    def forward(self, img):
        """
        Forward pass of the discriminator
        
        Args:
            img (torch.Tensor): Batch of images with shape (batch_size, channels, img_size, img_size)
            
        Returns:
            torch.Tensor: Classification scores, shape (batch_size, 1)
        """
        validity = self.model(img)
        validity = validity.view(validity.size(0), -1)
        return validity