import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from collections import OrderedDict
from utils import dummy_context_mgr


class CLIP_IMG_ENCODER(nn.Module):
    """
       CLIP_IMG_ENCODER module for encoding images using CLIP's visual transformer.
    """

    def __init__(self, CLIP):
        """
        Initialize the CLIP_IMG_ENCODER module.

        Args:
            CLIP (CLIP): Pre-trained CLIP model.
        """
        super(CLIP_IMG_ENCODER, self).__init__()
        model = CLIP.visual
        self.define_module(model)
        # freeze the parameters of the CLIP model
        for param in self.parameters():
            param.requires_grad = False

    def define_module(self, model):
        """
        Define the individual layers and modules of the CLIP visual transformer model.
        Args:
            model (nn.Module): CLIP visual transformer model.
        """
        # Extract required modules from the CLIP model
        self.conv1 = model.conv1  # Convolutional layer
        self.class_embedding = model.class_embedding  # Class embedding layer
        self.positional_embedding = model.positional_embedding  # Positional embedding layer
        self.ln_pre = model.ln_pre  # Linear Normalization layer for pre-normalization
        self.transformer = model.transformer  # Transformer block
        self.ln_post = model.ln_post  # Linear Normalization layer for post-normalization
        self.proj = model.proj  # projection matrix

    @property
    def dtype(self):
        """
         Get the data type of the convolutional layer weights.
        """
        return self.conv1.weight.dtype

    def transf_to_CLIP_input(self, inputs):
        """
        Transform input images to the format expected by CLIP.

        Args:
            inputs (torch.Tensor): Input images.

        Returns:
            torch.Tensor: Transformed images.
        """
        device = inputs.device
        # Check the size of the input image tensor
        if len(inputs.size()) != 4:
            raise ValueError('Expect the (B, C, X, Y) tensor.')
        else:
            # Normalize input images
            mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0).to(device)
            var = torch.tensor([0.26862954, 0.26130258, 0.27577711]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0).to(device)
            inputs = F.interpolate(inputs * 0.5 + 0.5, size=(224, 224))
            inputs = ((inputs + 1) * 0.5 - mean) / var
            return inputs

    def forward(self, img: torch.Tensor):
        """
        Forward pass of the CLIP_IMG_ENCODER module.

        Args:
            img (torch.Tensor): Input images.

        Returns:
            torch.Tensor: Local features extracted from the image.
            torch.Tensor: Encoded image embeddings.
        """
        # Transform input images to the format expected by CLIP and set its datatype appropriately
        x = self.transf_to_CLIP_input(img)
        x = x.type(self.dtype)

        # Pass the image through Convolutional layer
        x = self.conv1(x)  # shape = [*, width, grid, grid]
        grid = x.size(-1)

        # Reshape and permute the tensor for transformer input
        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        # Add class and positional embeddings
        x = torch.cat(
            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
        x = x + self.positional_embedding.to(x.dtype)
        x = self.ln_pre(x)

        # NLD (Batch Size - Length - Dimension) -> LND (Length - Batch Size - Dimension)
        x = x.permute(1, 0, 2)

        # Extract local features using transformer blocks
        selected = [1, 4, 8]
        local_features = []
        for i in range(12):
            x = self.transformer.resblocks[i](x)
            if i in selected:
                local_features.append(
                    x.permute(1, 0, 2)[:, 1:, :].permute(0, 2, 1).reshape(-1, 768, grid, grid).contiguous().type(
                        img.dtype))
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_post(x[:, 0, :])
        if self.proj is not None:
            x = x @ self.proj  # Perform matrix multiplication with projection matrix and tensor
        return torch.stack(local_features, dim=1), x.type(img.dtype)


class CLIP_TXT_ENCODER(nn.Module):
    """
        CLIP_TXT_ENCODER module for encoding text inputs using CLIP's transformer.
    """

    def __init__(self, CLIP):
        """
        Initialize the CLIP_TXT_ENCODER module.

        Args:
            CLIP (CLIP): Pre-trained CLIP model.
        """
        super(CLIP_TXT_ENCODER, self).__init__()
        self.define_module(CLIP)
        # Freeze the parameters of the CLIP model
        for param in self.parameters():
            param.requires_grad = False

    def define_module(self, CLIP):
        """
        Define the individual modules of the CLIP transformer model.

        Args:
            CLIP (CLIP): Pre-trained CLIP model.
        """
        self.transformer = CLIP.transformer  # Transformer block
        self.vocab_size = CLIP.vocab_size  # Size of the vocabulary of the transformer
        self.token_embedding = CLIP.token_embedding  # token embedding block
        self.positional_embedding = CLIP.positional_embedding  # positional embedding block
        self.ln_final = CLIP.ln_final  # Linear Normalization layer
        self.text_projection = CLIP.text_projection  # Projection matrix for text

    @property
    def dtype(self):
        """
        Get the data type of the first layer's weights in the transformer.
        """
        return self.transformer.resblocks[0].mlp.c_fc.weight.dtype

    def forward(self, text):
        """
        Forward pass of the CLIP_TXT_ENCODER module.

        Args:
            text (torch.Tensor): Input text tokens.

        Returns:
            torch.Tensor: Encoded sentence embeddings.
            torch.Tensor: Transformer output for the input text.
        """
        # Embed input text tokens
        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
        # Add positional embeddings
        x = x + self.positional_embedding.type(self.dtype)
        # Permute dimensions for transformer input
        x = x.permute(1, 0, 2)  # NLD -> LND
        # Pass input through the transformer
        x = self.transformer(x)
        # Permute dimensions back to original shape
        x = x.permute(1, 0, 2)  # LND -> NLD
        # Apply layer normalization
        x = self.ln_final(x).type(self.dtype)  # shape = [batch_size, n_ctx, transformer.width]
        # Extract sentence embeddings from the end-of-text (eot_token : is the highest number in each sequence)
        sent_emb = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection

        # Return the sentence embedding and transformer ouput
        return sent_emb, x


class CLIP_Mapper(nn.Module):
    """
    CLIP_Mapper module for mapping images with prompts using CLIP's transformer.
    """

    def __init__(self, CLIP):
        """
        Initialize the CLIP_Mapper module.

        Args:
            CLIP (CLIP): Pre-trained CLIP model.
        """
        super(CLIP_Mapper, self).__init__()
        model = CLIP.visual
        self.define_module(model)
        # Freeze the parameters of the CLIP visual model
        for param in model.parameters():
            param.requires_grad = False

    def define_module(self, model):
        """
        Define the individual modules of the CLIP visual model.

        Args:
            model: Pre-trained CLIP visual model.
        """
        self.conv1 = model.conv1
        self.class_embedding = model.class_embedding
        self.positional_embedding = model.positional_embedding
        self.ln_pre = model.ln_pre
        self.transformer = model.transformer

    @property
    def dtype(self):
        """
        Get the data type of the weights of the first convolutional layer.
        """
        return self.conv1.weight.dtype

    def forward(self, img: torch.Tensor, prompts: torch.Tensor):
        """
        Forward pass of the CLIP_Mapper module.

        Args:
            img (torch.Tensor): Input image tensor.
            prompts (torch.Tensor): Prompt tokens for mapping.

        Returns:
            torch.Tensor: Mapped features from the CLIP model.
        """

        # Convert input image and prompts to the appropriate data type
        x = img.type(self.dtype)
        prompts = prompts.type(self.dtype)
        grid = x.size(-1)

        # Reshape the input image tensor
        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        # Append the class embeddings to input tensors
        x = torch.cat(
            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
             x],
            dim=1
        )  # shape = [*, grid ** 2 + 1, width]

        # Append the positional embeddings to the input tensor
        x = x + self.positional_embedding.to(x.dtype)

        # Perform the layer normalization
        x = self.ln_pre(x)
        # NLD -> LND
        x = x.permute(1, 0, 2)
        # Local features
        selected = [1, 2, 3, 4, 5, 6, 7, 8]
        begin, end = 0, 12
        prompt_idx = 0
        for i in range(begin, end):
            # Add prompt to the input tensor
            if i in selected:
                prompt = prompts[:, prompt_idx, :].unsqueeze(0)
                prompt_idx = prompt_idx + 1
                x = torch.cat((x, prompt), dim=0)
                x = self.transformer.resblocks[i](x)
                x = x[:-1, :, :]
            else:
                x = self.transformer.resblocks[i](x)
        # Reshape and return mapped features
        return x.permute(1, 0, 2)[:, 1:, :].permute(0, 2, 1).reshape(-1, 768, grid, grid).contiguous().type(img.dtype)


class CLIP_Adapter(nn.Module):
    """
    CLIP_Adapter module for adapting features from a generator to match the CLIP model's input requirements.
    """

    def __init__(self, in_ch, mid_ch, out_ch, G_ch, CLIP_ch, cond_dim, k, s, p, map_num, CLIP):
        """
        Initialize the CLIP_Adapter module.

        Args:
            in_ch (int): Number of input channels.
            mid_ch (int): Number of channels in the intermediate layers.
            out_ch (int): Number of output channels.
            G_ch (int): Number of channels in the generator's output.
            CLIP_ch (int): Number of channels in the CLIP model's input.
            cond_dim (int): Dimension of the conditioning vector.
            k (int): Kernel size for convolutional layers.
            s (int): Stride for convolutional layers.
            p (int): Padding for convolutional layers.
            map_num (int): Number of mapping blocks.
            CLIP: Pre-trained CLIP model.
        """
        super(CLIP_Adapter, self).__init__()
        self.CLIP_ch = CLIP_ch
        self.FBlocks = nn.ModuleList([])
        # Define Mapping blocks (M_Block) and them to Feature blocks (FBlock) for given number of mapping blocks.
        self.FBlocks.append(M_Block(in_ch, mid_ch, out_ch, cond_dim, k, s, p))
        for i in range(map_num - 1):
            self.FBlocks.append(M_Block(out_ch, mid_ch, out_ch, cond_dim, k, s, p))
        # Convolutional layer to fuse adapted features
        self.conv_fuse = nn.Conv2d(out_ch, CLIP_ch, 5, 1, 2)
        # CLIP Mapper module to map adapted features to CLIP's input space
        self.CLIP_ViT = CLIP_Mapper(CLIP)
        # Convolutional layer to further process mapped features
        self.conv = nn.Conv2d(768, G_ch, 5, 1, 2)
        # Fully connected layer for conditioning
        self.fc_prompt = nn.Linear(cond_dim, CLIP_ch * 8)

    def forward(self, out, c):
        """
        Forward pass of the CLIP_Adapter module. Takes output features from the generator and conditioning vector
        as input, adapts features using the Feature block having multiple mapping blocks, fuses them, map them to
        CLIPs input space and returns the processed features

        Args:
            out (torch.Tensor): Output features from the generator.
            c (torch.Tensor): Conditioning vector.

        Returns:
            torch.Tensor: Adapted and mapped features for the generator.
        """

        # Generate prompts from the conditioning vector
        prompts = self.fc_prompt(c).view(c.size(0), -1, self.CLIP_ch)

        # Pass features through feature block consisting of multiple mapping blocks
        for FBlock in self.FBlocks:
            out = FBlock(out, c)
        # Fuse adapted features
        fuse_feat = self.conv_fuse(out)
        # Map fused features to CLIP's input space
        map_feat = self.CLIP_ViT(fuse_feat, prompts)
        # Further process mapped features and return
        return self.conv(fuse_feat + 0.1 * map_feat)


class NetG(nn.Module):
    """
    Generator network for synthesizing images conditioned on text and noise
    """

    def __init__(self, ngf, nz, cond_dim, imsize, ch_size, mixed_precision, CLIP):
        """
        Initializes the Generator network.

        Parameters:
            ngf (int): Number of generator filters.
            nz (int): Dimensionality of the input noise vector.
            cond_dim (int): Dimensionality of the conditioning vector.
            imsize (int): Size of the generated images.
            ch_size (int): Number of output channels for the generated images.
            mixed_precision (bool): Whether to use mixed precision training.
            CLIP: CLIP model for feature adaptation.

        """
        super(NetG, self).__init__()
        # Define attributes
        self.ngf = ngf
        self.mixed_precision = mixed_precision

        # Build CLIP Mapper
        self.code_sz, self.code_ch, self.mid_ch = 7, 64, 32
        self.CLIP_ch = 768
        # fully connected layer to convert the noise vector into a feature map of dimensions (code_sz * code_sz * code_ch)
        self.fc_code = nn.Linear(nz, self.code_sz * self.code_sz * self.code_ch)
        self.mapping = CLIP_Adapter(self.code_ch, self.mid_ch, self.code_ch, ngf * 8, self.CLIP_ch, cond_dim + nz, 3, 1,
                                    1, 4, CLIP)
        # Build GBlocks
        self.GBlocks = nn.ModuleList([])
        in_out_pairs = list(get_G_in_out_chs(ngf, imsize))
        imsize = 4
        for idx, (in_ch, out_ch) in enumerate(in_out_pairs):
            if idx < (len(in_out_pairs) - 1):
                imsize = imsize * 2
            else:
                imsize = 224
            self.GBlocks.append(G_Block(cond_dim + nz, in_ch, out_ch, imsize))

        # To RGB image conversion using the sequential layers having leakyReLU activation function
        self.to_rgb = nn.Sequential(
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(out_ch, ch_size, 3, 1, 1),
        )

    def forward(self, noise, c, eval=False):  # x=noise, c=ent_emb
        """
        Forward pass of the generator network.

        Args:
            noise (torch.Tensor): Input noise vector.
            c (torch.Tensor): Conditioning information, typically an embedding representing attributes of the output.
            eval (bool, optional): Flag indicating whether the network is in evaluation mode. Defaults to False.

        Returns:
            torch.Tensor: Generated RGB images.
        """
        # Context manager for enabling automatic mixed precision training
        with torch.cuda.amp.autocast() if self.mixed_precision and not eval else dummy_context_mgr() as mp:
            # Concatenate noise and conditioning information
            cond = torch.cat((noise, c), dim=1)

            # Pass noise through fully connected layer to generate feature map and adapt features using CLIP Mapper
            out = self.mapping(self.fc_code(noise).view(noise.size(0), self.code_ch, self.code_sz, self.code_sz), cond)

            # Apply GBlocks to progressively upsample feature representation, fuse text and visual features
            for GBlock in self.GBlocks:
                out = GBlock(out, cond)

            # Convert final feature representation to RGB images
            out = self.to_rgb(out)

        return out


class NetD(nn.Module):
    """
    Discriminator network for evaluating the realism of images.
    Attributes:
        DBlocks (nn.ModuleList): List of D_Block modules for processing feature maps.
        main (D_Block): Main D_Block module for final processing.
    """

    def __init__(self, ndf, imsize, ch_size, mixed_precision):
        """
        Initializes the Discriminator network

        Args:
        ndf (int): Number of channels in the initial features.
        imsize (int): Size of the input images (assumed square).
        ch_size (int): Number of channels in the output feature maps.
        mixed_precision (bool): Flag indicating whether to use mixed precision training.
        """
        super(NetD, self).__init__()
        self.mixed_precision = mixed_precision
        # Define the DBlock
        self.DBlocks = nn.ModuleList([
            D_Block(768, 768, 3, 1, 1, res=True, CLIP_feat=True),
            D_Block(768, 768, 3, 1, 1, res=True, CLIP_feat=True),
        ])
        # Define the main DBlock for the final processing
        self.main = D_Block(768, 512, 3, 1, 1, res=True, CLIP_feat=False)

    def forward(self, h):
        """
        Forward pass of the discriminator network.
        Args:
            h (torch.Tensor): Input feature maps.
        Returns:
            torch.Tensor: Discriminator output.
        """
        with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
            # Initial feature map
            out = h[:, 0]
            # Pass the input feature through each DBlock
            for idx in range(len(self.DBlocks)):
                out = self.DBlocks[idx](out, h[:, idx + 1])
            # Final processing through the main DBlock
            out = self.main(out)
        return out


class NetC(nn.Module):
    """
    Classifier / Comparator network for classifying the joint features of the generator output and condition text.
    Attributes:
        cond_dim (int): Dimensionality of the conditioning information.
        mixed_precision (bool): Flag indicating whether to use mixed precision training.
        joint_conv (nn.Sequential): Sequential module defining the classifier layers.
    """
    def __init__(self, ndf, cond_dim, mixed_precision):
        """

        """
        super(NetC, self).__init__()
        self.cond_dim = cond_dim
        self.mixed_precision = mixed_precision
        # Define the classifier layers, sequential convolutional 2D layer with LeakyReLU as the activation function
        self.joint_conv = nn.Sequential(
            nn.Conv2d(512 + 512, 128, 4, 1, 0, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 1, 4, 1, 0, bias=False),
        )

    def forward(self, out, cond):
        """
        Forward pass of the classifier network.

        Args:
            out (torch.Tensor): Generator output feature map.
            cond (torch.Tensor): Conditioning information vector
        """
        with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
            # Reshape and repeat conditioning information vector to match the feature map size
            cond = cond.view(-1, self.cond_dim, 1, 1)
            cond = cond.repeat(1, 1, 7, 7)

            # Concatenate feature map and conditioned information
            h_c_code = torch.cat((out, cond), 1)

            # Pass through the classifier layers
            out = self.joint_conv(h_c_code)
        return out


class M_Block(nn.Module):
    """
    Multi-scale block consisting of convolutional layers and conditioning.

    Attributes:
        conv1 (nn.Conv2d): First convolutional layer.
        fuse1 (DFBlock): Conditioning block for the first convolutional layer.
        conv2 (nn.Conv2d): Second convolutional layer.
        fuse2 (DFBlock): Conditioning block for the second convolutional layer.
        learnable_sc (bool): Flag indicating whether the shortcut connection is learnable.
        c_sc (nn.Conv2d): Convolutional layer for the shortcut connection.

    """
    def __init__(self, in_ch, mid_ch, out_ch, cond_dim, k, s, p):
        """
        Initializes the Multi-scale block.

        Args:
            in_ch (int): Number of input channels.
            mid_ch (int): Number of channels in the intermediate layers.
            out_ch (int): Number of output channels.
            cond_dim (int): Dimensionality of the conditioning information.
            k (int): Kernel size for convolutional layers.
            s (int): Stride for convolutional layers.
            p (int): Padding for convolutional layers.

        """
        super(M_Block, self).__init__()

        # Define convolutional layers and conditioning blocks
        self.conv1 = nn.Conv2d(in_ch, mid_ch, k, s, p)
        self.fuse1 = DFBLK(cond_dim, mid_ch)
        self.conv2 = nn.Conv2d(mid_ch, out_ch, k, s, p)
        self.fuse2 = DFBLK(cond_dim, out_ch)

        # Learnable shortcut connection
        self.learnable_sc = in_ch != out_ch
        if self.learnable_sc:
            self.c_sc = nn.Conv2d(in_ch, out_ch, 1, stride=1, padding=0)

    def shortcut(self, x):
        """
        Defines the shortcut connection.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Shortcut connection output.
        """
        if self.learnable_sc:
            x = self.c_sc(x)
        return x

    def residual(self, h, text):
        """
        Defines the residual path with conditioning.

        Args:
            h (torch.Tensor): Input tensor.
            text (torch.Tensor): Conditioning information.

        Returns:
            torch.Tensor: Residual path output.
        """
        h = self.conv1(h)
        h = self.fuse1(h, text)
        h = self.conv2(h)
        h = self.fuse2(h, text)
        return h

    def forward(self, h, c):
        """
        Forward pass of the multi-scale block.

        Args:
            h (torch.Tensor): Input tensor.
            c (torch.Tensor): Conditioning information.

        Returns:
            torch.Tensor: Output tensor.
        """
        return self.shortcut(h) + self.residual(h, c)


class G_Block(nn.Module):
    """
        Generator block consisting of convolutional layers and conditioning.

        Attributes:
            imsize (int): Size of the output image.
            learnable_sc (bool): Flag indicating whether the shortcut connection is learnable.
            c1 (nn.Conv2d): First convolutional layer.
            c2 (nn.Conv2d): Second convolutional layer.
            fuse1 (DFBLK): Conditioning block for the first convolutional layer.
            fuse2 (DFBLK): Conditioning block for the second convolutional layer.
            c_sc (nn.Conv2d): Convolutional layer for the shortcut connection.
        """

    def __init__(self, cond_dim, in_ch, out_ch, imsize):
        """
        Initialize the Generator block.

        Args:
            cond_dim (int): Dimensionality of the conditioning information.
            in_ch (int): Number of input channels.
            out_ch (int): Number of output channels.
            imsize (int): Size of the output image.
        """
        super(G_Block, self).__init__()

        # Initialize attributes
        self.imsize = imsize
        self.learnable_sc = in_ch != out_ch

        # Define convolution layers and conditioning blocks
        self.c1 = nn.Conv2d(in_ch, out_ch, 3, 1, 1)
        self.c2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1)
        self.fuse1 = DFBLK(cond_dim, in_ch)
        self.fuse2 = DFBLK(cond_dim, out_ch)

        # Learnable shortcut connection
        if self.learnable_sc:
            self.c_sc = nn.Conv2d(in_ch, out_ch, 1, stride=1, padding=0)

    def shortcut(self, x):
        """
        Defines the shortcut connection.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Shortcut connection output.
        """
        if self.learnable_sc:
            x = self.c_sc(x)
        return x

    def residual(self, h, y):
        """
        Defines the residual path with conditioning.

        Args:
            h (torch.Tensor): Input tensor.
            y (torch.Tensor): Conditioning information.

        Returns:
            torch.Tensor: Residual path output.
        """
        h = self.fuse1(h, y)
        h = self.c1(h)
        h = self.fuse2(h, y)
        h = self.c2(h)
        return h

    def forward(self, h, y):
        """
        Forward pass of the generator block.

        Args:
            h (torch.Tensor): Input tensor.
            y (torch.Tensor): Conditioning information.

        Returns:
            torch.Tensor: Output tensor.
        """
        h = F.interpolate(h, size=(self.imsize, self.imsize))
        return self.shortcut(h) + self.residual(h, y)


class D_Block(nn.Module):
    """
    Discriminator block.
    """
    def __init__(self, fin, fout, k, s, p, res, CLIP_feat):
        """
        Initializes Discriminator block.

        Args:
        - fin (int): Number of input channels.
        - fout (int): Number of output channels.
        - k (int): Kernel size for convolutional layers.
        - s (int): Stride for convolutional layers.
        - p (int): Padding for convolutional layers.
        - res (bool): Whether to use residual connection.
        - CLIP_feat (bool): Whether to incorporate CLIP features.
        """
        super(D_Block, self).__init__()
        self.res, self.CLIP_feat = res, CLIP_feat
        self.learned_shortcut = (fin != fout)

        # Convolutional layers for residual path
        self.conv_r = nn.Sequential(
            nn.Conv2d(fin, fout, k, s, p, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(fout, fout, k, s, p, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
        )

        # Convolutional layers for shortcut connection
        self.conv_s = nn.Conv2d(fin, fout, 1, stride=1, padding=0)

        # Parameters for learned residual and CLIP features
        if self.res == True:
            self.gamma = nn.Parameter(torch.zeros(1))
        if self.CLIP_feat == True:
            self.beta = nn.Parameter(torch.zeros(1))

    def forward(self, x, CLIP_feat=None):
        """
        Forward pass of the discriminator block.

        Args:
        - x (torch.Tensor): Input tensor.
        - CLIP_feat (torch.Tensor): Optional CLIP features tensor.

        Returns:
        - torch.Tensor: Output tensor.
        """
        # Compute the residual features
        res = self.conv_r(x)

        # Compute the shortcut connection
        if self.learned_shortcut:
            x = self.conv_s(x)

        # Incorporate learned residual and CLIP features if enabled
        if (self.res == True) and (self.CLIP_feat == True):
            return x + self.gamma * res + self.beta * CLIP_feat
        elif (self.res == True) and (self.CLIP_feat != True):
            return x + self.gamma * res
        elif (self.res != True) and (self.CLIP_feat == True):
            return x + self.beta * CLIP_feat
        else:
            return x


class DFBLK(nn.Module):
    """
    Diffusion Block of the Generator network with Conditional feature block
    """
    def __init__(self, cond_dim, in_ch):
        """
        Initializing the Conditional feature block of the DFBlock.

        Args:
        - cond_dim (int): Dimensionality of the conditional input.
        - in_ch (int): Number of input channels.
        """
        super(DFBLK, self).__init__()
        # Define conditional affine transformations
        self.affine0 = Affine(cond_dim, in_ch)
        self.affine1 = Affine(cond_dim, in_ch)

    def forward(self, x, y=None):
        """
        Forward pass of the conditional feature block.

        Args:
        - x (torch.Tensor): Input tensor.
        - y (torch.Tensor, optional): Conditional input tensor. Default is None.

        Returns:
        - torch.Tensor: Output tensor.
        """
        # Apply the first affine transformation and activation function
        h = self.affine0(x, y)
        h = nn.LeakyReLU(0.2, inplace=True)(h)
        # Apply second affine transformation and activation function
        h = self.affine1(h, y)
        h = nn.LeakyReLU(0.2, inplace=True)(h)
        return h


class QuickGELU(nn.Module):
    """
    Efficient and faster version of GELU,
    for non-linearity and to learn complex patterns
    """
    def forward(self, x: torch.Tensor):
        """
        Forward pass of the QuickGELU activation function.

        Args:
        - x (torch.Tensor): Input tensor.

        Returns:
        - torch.Tensor: Output tensor.
        """
        # Apply QuickGELU activation function
        return x * torch.sigmoid(1.702 * x)


# Taken from the RAT-GAN repository
class Affine(nn.Module):
    """
    Affine transformation module that applies conditional scaling and shifting to input features,
    to incorporate additional control over the generated output based on input conditions.
    """
    def __init__(self, cond_dim, num_features):
        """
        Initialize the affine transformation module.
        Args:
            cond_dim (int): Dimensionality of the conditioning information.
            num_features (int): Number of input features.
        """
        super(Affine, self).__init__()
        # Define 2 fully connected networks to compute gamma and beta parameters
        # each 2 linear layers with RELU activation in between
        self.fc_gamma = nn.Sequential(OrderedDict([
            ('linear1', nn.Linear(cond_dim, num_features)),
            ('relu1', nn.ReLU(inplace=True)),
            ('linear2', nn.Linear(num_features, num_features)),
        ]))
        self.fc_beta = nn.Sequential(OrderedDict([
            ('linear1', nn.Linear(cond_dim, num_features)),
            ('relu1', nn.ReLU(inplace=True)),
            ('linear2', nn.Linear(num_features, num_features)),
        ]))
        # Initializes the weights and biases of the network
        self._initialize()

    def _initialize(self):
        """
        Initializes the weights and biases of the linear layers responsible for computing gamma and beta
        """
        nn.init.zeros_(self.fc_gamma.linear2.weight.data)
        nn.init.ones_(self.fc_gamma.linear2.bias.data)
        nn.init.zeros_(self.fc_beta.linear2.weight.data)
        nn.init.zeros_(self.fc_beta.linear2.bias.data)

    def forward(self, x, y=None):
        """
        Forward pass of the Affine transformation module.

        Args:
            x (torch.Tensor): Input tensor.
            y (torch.Tensor, optional): Conditioning information tensor. Default is None.

        Returns:
            torch.Tensor: Transformed tensor after applying affine transformation.
        """
        # Compute gamma and beta parameters
        weight = self.fc_gamma(y)
        bias = self.fc_beta(y)

        # Ensure proper shape for weight and bias tensors
        if weight.dim() == 1:
            weight = weight.unsqueeze(0)
        if bias.dim() == 1:
            bias = bias.unsqueeze(0)

        # Expand weight and bias tensors to match input tensor shape
        size = x.size()
        weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
        bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)

        # Apply affine transformation
        return weight * x + bias


def get_G_in_out_chs(nf, imsize):
    """
    Compute input-output channel pairs for generator blocks based on given number of channels and image size.

    Args:
        nf (int): Number of input channels.
        imsize (int): Size of the input image.

    Returns:
        list: List of tuples containing input-output channel pairs for generator blocks.
    """
    # Determine the number of layers based on image size
    layer_num = int(np.log2(imsize)) - 1

    # Compute the number of channels for each layer
    channel_nums = [nf * min(2 ** idx, 8) for idx in range(layer_num)]

    # Reverse the channel numbers to start with the highest channel count
    channel_nums = channel_nums[::-1]

    # Generate input-output channel pairs for generator blocks
    in_out_pairs = zip(channel_nums[:-1], channel_nums[1:])

    return in_out_pairs


def get_D_in_out_chs(nf, imsize):
    """
    Compute input-output channel pairs for discriminator blocks based on given number of channels and image size.

    Args:
        nf (int): Number of input channels.
        imsize (int): Size of the input image.

    Returns:
        list: List of tuples containing input-output channel pairs for discriminator blocks.
    """
    # Determine the number of layers based on image size
    layer_num = int(np.log2(imsize)) - 1

    # Compute the number of channels for each layer
    channel_nums = [nf * min(2 ** idx, 8) for idx in range(layer_num)]

    # Generate input-output channel pairs for discriminator blocks
    in_out_pairs = zip(channel_nums[:-1], channel_nums[1:])

    return in_out_pairs