"""
Contains the CNN model for malware classification.
"""

import torch.nn as nn


class MalwareCNN(nn.Module):
    """
    Convolutional Neural Network for classifying malware families based on byte images.

    Args:
        num_classes (int): Number of unique malware families to classify.
    """

    def __init__(self, num_classes=24):
        super(MalwareCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 128 -> 64
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 64 -> 32
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 32 -> 16
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 16 * 16, 512),
            nn.ReLU(),
            nn.Dropout(0.6),  # Increased dropout to prevent overfitting
            nn.Linear(512, num_classes),
        )

    def forward(self, x):
        """
        Forward pass through the CNN.

        Args:
            x (torch.Tensor): Input batch of images.

        Returns:
            torch.Tensor: Logits for each class.
        """
        x = self.features(x)
        x = self.classifier(x)
        return x