we have append below logic techniques in the code 


1. **Random Resizing and Padding**
   \[
   x_{\text{resized}} = \text{resize}(x, s_{\text{new}}), \quad x_{\text{padded}} = \text{pad}(x_{\text{resized}}, p)
   \]
   where \( s_{\text{new}} \) is a randomly selected size and \( p \) is padding added to preserve input dimensions.

2. **PGD Adversarial Training**
   \[
   \theta = \arg \min_{\theta} \mathbb{E}_{(x, y) \sim D} \left[ \max_{\|\delta\| \leq \epsilon} L(f_\theta(x + \delta), y) \right]
   \]
   where \( \delta \) is adversarial perturbation and \( \epsilon \) defines the perturbation bound.

3. **Fourier Transformation Defense**
   \[
   X_{\text{fft}} = \mathcal{F}(x), \quad x_{\text{filtered}} = \mathcal{F}^{-1}(X_{\text{fft}})
   \]
   where \( \mathcal{F} \) and \( \mathcal{F}^{-1} \) are Fourier transform and inverse.

4. **Gradient Masking**
   \[
   x_{\text{masked}} = x - \eta \nabla_x L(x, y)
   \]
   where \( \eta \) is the masking strength.

5. **Adversarial Logit Pairing**
   \[
   \mathcal{L}_{\text{ALP}} = L(f(x), y) + \alpha \cdot \|f(x) - f(x_{\text{adv}})\|^2
   \]
   where \( f(x_{\text{adv}}) \) is the logit of adversarial examples, and \( \alpha \) controls regularization.

6. **Spectral Normalization**
   \[
   W_{\text{norm}} = \frac{W}{\sigma(W)}
   \]
   where \( \sigma(W) \) is the largest singular value of weight \( W \).

7. **Autoencoder Defense**
   \[
   x_{\text{denoised}} = \text{decoder}(\text{encoder}(x))
   \]
   using encoder-decoder structure to remove perturbations.

8. **Self-Ensemble Adversarial Training**
   \[
   \theta_{\text{new}} = \theta_{\text{old}} + \eta \cdot \frac{\partial L(f_\theta(x + \delta), y)}{\partial \theta}
   \]
   updating model with perturbation ensemble.

9. **Contextual Embedding Smoothing**
   \[
   x_{\text{smoothed}} = x + \lambda \cdot \nabla_x L(f_\theta(x), y)
   \]
   smoothing out context features.

10. **Perceptual Loss Minimization**
   \[
   \mathcal{L}_{\text{perceptual}} = \|f(x) - f(x_{\text{adv}})\|_{p}
   \]
   enforcing similarity in feature space.
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from typing import Tuple, Optional, Any
import math
import random
import warnings


class AdversarialDefense:
    """
    A class encapsulating various adversarial defense techniques.
    """

    def __init__(self, model: nn.Module, loss_fn: nn.Module, optimizer: optim.Optimizer) -> None:
        """
        Initializes the AdversarialDefense with a model, loss function, and optimizer.

        Args:
            model (nn.Module): The neural network model.
            loss_fn (nn.Module): The loss function.
            optimizer (optim.Optimizer): The optimizer.
        """
        if not isinstance(model, nn.Module):
            raise TypeError("model must be an instance of nn.Module")
        if not isinstance(loss_fn, nn.Module):
            raise TypeError("loss_fn must be an instance of nn.Module")
        if not isinstance(optimizer, optim.Optimizer):
            raise TypeError("optimizer must be an instance of optim.Optimizer")

        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    def cutmix(
        self,
        images: torch.Tensor,
        labels: torch.Tensor,
        alpha: float = 1.0
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Applies the CutMix augmentation technique.

        Args:
            images (torch.Tensor): Batch of images with shape (batch_size, channels, height, width).
            labels (torch.Tensor): Corresponding labels with shape (batch_size, num_classes).
            alpha (float): Hyperparameter for the Beta distribution.

        Returns:
            Tuple containing:
                - mixed_images (torch.Tensor): Batch of mixed images.
                - target_a (torch.Tensor): Original labels.
                - target_b (torch.Tensor): Shuffled labels.
                - lam (float): Lambda value used for mixing.
        """
        if images.size(0) != labels.size(0):
            raise ValueError("Batch size of images and labels must be the same")

        lam = np.random.beta(alpha, alpha) if alpha > 0 else 1
        batch_size, _, height, width = images.size()
        cx, cy, w, h = self._rand_bbox(width, height, lam)

        images_shuffled = self._shuffle_batch(images)
        labels_shuffled = self._shuffle_batch(labels)

        mixed_images = images.clone()
        mixed_images[:, :, cy:cy + h, cx:cx + w] = images_shuffled[:, :, cy:cy + h, cx:cx + w]
        lam = 1 - ((w * h) / (width * height))

        return mixed_images, labels, labels_shuffled, lam

    def cutout(
        self,
        images: torch.Tensor,
        mask_size: int = 16
    ) -> torch.Tensor:
        """
        Applies the Cutout augmentation technique.

        Args:
            images (torch.Tensor): Batch of images with shape (batch_size, channels, height, width).
            mask_size (int): Size of the square mask.

        Returns:
            masked_images (torch.Tensor): Batch of masked images.
        """
        if mask_size <= 0:
            raise ValueError("mask_size must be positive")

        batch_size, channels, height, width = images.size()
        masked_images = images.clone()

        for i in range(batch_size):
            y = random.randint(0, height)
            x = random.randint(0, width)

            y1 = max(0, y - mask_size // 2)
            y2 = min(height, y + mask_size // 2)
            x1 = max(0, x - mask_size // 2)
            x2 = min(width, x + mask_size // 2)

            masked_images[i, :, y1:y2, x1:x2] = 0

        return masked_images

    def feature_squeezing(
        self,
        images: torch.Tensor,
        bits: int = 4
    ) -> torch.Tensor:
        """
        Applies the Feature Squeezing technique by reducing the precision of input features.

        Args:
            images (torch.Tensor): Batch of images.
            bits (int): Number of bits to keep.

        Returns:
            squeezed_images (torch.Tensor): Batch of squeezed images.
        """
        if not (1 <= bits <= 8):
            raise ValueError("bits must be in [1, 8]")

        squeezed_images = torch.round(images * (2 ** bits - 1)) / (2 ** bits - 1)
        return squeezed_images

    def gaussian_augmentation(
        self,
        images: torch.Tensor,
        sigma: float = 0.1
    ) -> torch.Tensor:
        """
        Applies Gaussian noise augmentation.

        Args:
            images (torch.Tensor): Batch of images.
            sigma (float): Standard deviation of the Gaussian noise.

        Returns:
            augmented_images (torch.Tensor): Batch of augmented images.
        """
        if sigma < 0:
            raise ValueError("sigma must be non-negative")

        noise = torch.randn_like(images) * sigma
        augmented_images = images + noise
        augmented_images = torch.clamp(augmented_images, 0.0, 1.0)
        return augmented_images

    def label_smoothing(
        self,
        labels: torch.Tensor,
        alpha: float = 0.1
    ) -> torch.Tensor:
        """
        Applies label smoothing.

        Args:
            labels (torch.Tensor): One-hot encoded labels.
            alpha (float): Smoothing parameter.

        Returns:
            smoothed_labels (torch.Tensor): Smoothed labels.
        """
        if not (0.0 <= alpha <= 1.0):
            raise ValueError("alpha must be in [0, 1]")

        num_classes = labels.size(1)
        smoothed_labels = labels * (1.0 - alpha) + alpha / num_classes
        return smoothed_labels

    def mixup(
        self,
        images: torch.Tensor,
        labels: torch.Tensor,
        alpha: float = 1.0
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
        """
        Applies the Mixup augmentation technique.

        Args:
            images (torch.Tensor): Batch of images.
            labels (torch.Tensor): Batch of labels.
            alpha (float): Parameter for the Beta distribution.

        Returns:
            Tuple containing:
                - mixed_images (torch.Tensor): Batch of mixed images.
                - target_a (torch.Tensor): Original labels.
                - target_b (torch.Tensor): Shuffled labels.
                - lam (float): Lambda value used for mixing.
        """
        if alpha <= 0:
            lam = 1
        else:
            lam = np.random.beta(alpha, alpha)

        batch_size = images.size(0)
        index = torch.randperm(batch_size).to(images.device)

        mixed_images = lam * images + (1 - lam) * images[index, :]
        target_a, target_b = labels, labels[index]
        return mixed_images, target_a, target_b, lam

    def spatial_smoothing(
        self,
        images: torch.Tensor,
        kernel_size: int = 3
    ) -> torch.Tensor:
        """
        Applies spatial smoothing to reduce high-frequency noise.

        Args:
            images (torch.Tensor): Batch of images.
            kernel_size (int): Size of the smoothing kernel.

        Returns:
            smoothed_images (torch.Tensor): Batch of smoothed images.
        """
        if kernel_size <= 0 or kernel_size % 2 == 0:
            raise ValueError("kernel_size must be a positive odd integer")

        padding = kernel_size // 2
        kernel = torch.ones((1, 1, kernel_size, kernel_size)) / (kernel_size ** 2)
        kernel = kernel.to(images.device)

        channels = images.size(1)
        kernel = kernel.expand(channels, 1, kernel_size, kernel_size)

        smoothed_images = F.conv2d(images, kernel, padding=padding, groups=channels)
        return smoothed_images

    def thermometer_encoding(
        self,
        images: torch.Tensor,
        threshold_levels: int = 10
    ) -> torch.Tensor:
        """
        Applies thermometer encoding to convert each feature into a binary vector.

        Args:
            images (torch.Tensor): Batch of images.
            threshold_levels (int): Number of threshold levels.

        Returns:
            encoded_images (torch.Tensor): Batch of thermometer encoded images.
        """
        if threshold_levels <= 0:
            raise ValueError("threshold_levels must be positive")

        max_val = images.max()
        thresholds = torch.linspace(0, max_val, steps=threshold_levels).to(images.device)
        encoded = (images.unsqueeze(1) >= thresholds.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)).float()
        return encoded

    def thermometer_encoding(
        self,
        images: torch.Tensor,
        threshold_levels: int = 10
        ) -> torch.Tensor:
        """
        Applies thermometer encoding to convert each feature into a binary vector.
    
        Args:
            images (torch.Tensor): Batch of images.
            threshold_levels (int): Number of threshold levels.
    
        Returns:
            encoded_images (torch.Tensor): Batch of thermometer encoded images.
                                            Shape: [batch_size, channels * threshold_levels, height, width]
        """
        if threshold_levels <= 0:
            raise ValueError("threshold_levels must be positive")
    
        max_val = images.max()
        thresholds = torch.linspace(0, max_val, steps=threshold_levels).to(images.device)
        
        # Reshape images to [batch_size, channels, 1, height, width]
        images_expanded = images.unsqueeze(2)
        
        # Reshape thresholds to [1, 1, threshold_levels, 1, 1]
        thresholds_expanded = thresholds.view(1, 1, threshold_levels, 1, 1)
        
        # Perform the comparison
        encoded = (images_expanded >= thresholds_expanded).float()  # Shape: [batch_size, channels, threshold_levels, height, width]
        
        # Optionally, reshape to merge channels and threshold levels
        encoded = encoded.view(images.size(0), images.size(1) * threshold_levels, images.size(2), images.size(3))
        
        return encoded
    def _rand_bbox(self, width: int, height: int, lam: float) -> Tuple[int, int, int, int]:
        """
        Generates a random bounding box for CutMix.

        Args:
            width (int): Width of the image.
            height (int): Height of the image.
            lam (float): Lambda value.

        Returns:
            Tuple containing (cx, cy, w, h): center coordinates and size of the box.
        """
        r_x = random.uniform(0, width)
        r_y = random.uniform(0, height)
        r_w = width * math.sqrt(1 - lam)
        r_h = height * math.sqrt(1 - lam)

        cx = int(r_x - r_w / 2)
        cy = int(r_y - r_h / 2)
        w = int(r_w)
        h = int(r_h)

        return cx, cy, w, h

    def _shuffle_batch(self, tensor: torch.Tensor) -> torch.Tensor:
        """
        Shuffles a batch of tensors.

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

        Returns:
            torch.Tensor: Shuffled tensor.
        """
        indices = torch.randperm(tensor.size(0)).to(tensor.device)
        return tensor[indices]

    def _total_variation_loss(self, images: torch.Tensor) -> torch.Tensor:
        """
        Computes the total variation loss.

        Args:
            images (torch.Tensor): Batch of images.

        Returns:
            tv_loss (torch.Tensor): Total variation loss.
        """
        tv_h = torch.mean(torch.abs(images[:, :, 1:, :] - images[:, :, :-1, :]))
        tv_w = torch.mean(torch.abs(images[:, :, :, 1:] - images[:, :, :, :-1]))
        return tv_h + tv_w

    # Placeholder methods for defenses requiring external modules
    def defense_gan(
        self,
        images: torch.Tensor
    ) -> torch.Tensor:
        """
        Projects inputs onto a GAN-generated manifold.
        Note: Requires a pre-trained GAN generator.

        Args:
            images (torch.Tensor): Batch of images.

        Returns:
            projected_images (torch.Tensor): Batch of projected images.
        """
        raise NotImplementedError("DefenseGAN requires a pre-trained GAN generator.")

    def jpeg_compression(
        self,
        images: torch.Tensor,
        quality: int = 75
    ) -> torch.Tensor:
        """
        Applies JPEG compression to images.
        Note: Requires external libraries such as PIL or torchvision.

        Args:
            images (torch.Tensor): Batch of images.
            quality (int): Quality factor for JPEG compression.

        Returns:
            compressed_images (torch.Tensor): Batch of JPEG-compressed images.
        """
        raise NotImplementedError("JPEG Compression requires external libraries.")

    
    def resample(
        self,
        images: torch.Tensor,
        new_size: Tuple[int, int]
    ) -> torch.Tensor:
        """
        Resamples images to a new resolution.

        Args:
            images (torch.Tensor): Batch of images.
            new_size (Tuple[int, int]): New size as (height, width).

        Returns:
            resampled_images (torch.Tensor): Batch of resampled images.
        """
        if not isinstance(new_size, tuple) or len(new_size) != 2:
            raise ValueError("new_size must be a tuple of (height, width)")

        resampled_images = F.interpolate(images, size=new_size, mode='bilinear', align_corners=False)
        return resampled_images


# Example Usage
if __name__ == "__main__":
    import numpy as np

    # Define a simple model for demonstration
    class SimpleCNN(nn.Module):
        def __init__(self, num_classes: int = 10) -> None:
            super(SimpleCNN, self).__init__()
            self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
            self.pool = nn.MaxPool2d(2, 2)
            self.fc1 = nn.Linear(16 * 16 * 16, num_classes)

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            x = self.pool(F.relu(self.conv1(x)))
            x = x.view(x.size(0), -1)
            x = self.fc1(x)
            return x

    # Initialize model, loss function, and optimizer
    model = SimpleCNN(num_classes=10)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # Initialize AdversarialDefense
    defense = AdversarialDefense(model, loss_fn, optimizer)

    # Dummy data
    batch_size = 8
    channels = 3
    height = 32
    width = 32
    num_classes = 10

    adversarial_samples = torch.rand(batch_size, channels, height, width)
    labels = F.one_hot(torch.randint(0, num_classes, (batch_size,)), num_classes=num_classes).float()

    # Apply CutMix
    mixed_images, target_a, target_b, lam = defense.cutmix(adversarial_samples, labels, alpha=1.0)
    print("CutMix Applied:")
    print(f"Lambda: {lam}")
    print(f"Mixed Images Shape: {mixed_images.shape}")

    # Apply Cutout
    cutout_images = defense.cutout(adversarial_samples, mask_size=8)
    print("\nCutout Applied:")
    print(f"Cutout Images Shape: {cutout_images.shape}")

    # Apply Feature Squeezing
    squeezed_images = defense.feature_squeezing(adversarial_samples, bits=4)
    print("\nFeature Squeezing Applied:")
    print(f"Squeezed Images Shape: {squeezed_images.shape}")

    # Apply Gaussian Augmentation
    gaussian_images = defense.gaussian_augmentation(adversarial_samples, sigma=0.05)
    print("\nGaussian Augmentation Applied:")
    print(f"Gaussian Images Shape: {gaussian_images.shape}")

    # Apply Label Smoothing
    smoothed_labels = defense.label_smoothing(labels, alpha=0.1)
    print("\nLabel Smoothing Applied:")
    print(f"Smoothed Labels: {smoothed_labels}")

    # Apply Mixup
    mixed_images_mixup, target_a_mixup, target_b_mixup, lam_mixup = defense.mixup(adversarial_samples, labels, alpha=1.0)
    print("\nMixup Applied:")
    print(f"Lambda: {lam_mixup}")
    print(f"Mixed Images Shape: {mixed_images_mixup.shape}")

    # Apply Spatial Smoothing
    smoothed_images = defense.spatial_smoothing(adversarial_samples, kernel_size=3)
    print("\nSpatial Smoothing Applied:")
    print(f"Smoothed Images Shape: {smoothed_images.shape}")

    # Apply Thermometer Encoding
    encoded_images = defense.thermometer_encoding(adversarial_samples, threshold_levels=10)
    print("\nThermometer Encoding Applied:")
    print(f"Encoded Images Shape: {encoded_images.shape}")