update enviroment with 

    metrics_calc_multi = MetricsCalculator(preds=model_output.logits , labels=labels, num_classes=200)
    metrics_multi = metrics_calc_multi.get_all_metrics()
    print("\nMulti-Class Classification Metrics:")
    print_metrics(metrics_multi)
# environment.py

import logging
from typing import Any, Dict, Optional, Tuple, List

import gym
from gym import spaces
import torch
import numpy as np
from torchvision import transforms

from attacks import AdversarialAttack, AttackConfig
from defence import (
    GaussianNoiseDefense,
    LabelSmoothingDefense,
    MixupDefense,
    FeatureSqueezingDefense,
    SpatialSmoothingDefense,
    CutoutDefense,
    JpegCompressionDefense,
)
from modelloader import ImageClassificationPipeline, ModelSource
from rlpedata import datapipeline, DatasetConfig
from Imageclassificationmetrics import MetricsCalculator

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AdversarialEnv(gym.Env):
    """
    OpenAI Gym environment to simulate adversarial attack scenarios and evaluate defense mechanisms.
    """
    
    metadata = {'render.modes': ['human']}

    def __init__(self, config: DatasetConfig, model_pipeline: ImageClassificationPipeline, attack_config: Optional[Dict[str, Any]] = None) -> None:
        """
        Initialize the adversarial environment.

        Args:
            config (DatasetConfig): Configuration for the dataset.
            model_pipeline (ImageClassificationPipeline): Pipeline to load the model.
            attack_config (Optional[Dict[str, Any]]): Configuration for attacks.
        """
        super(AdversarialEnv, self).__init__()

        self.config = config
        self.model_pipeline = model_pipeline
        self.model = self.model_pipeline.model.to(self.model_pipeline.device)
        self.model.eval()
        self.adversary = AdversarialAttack(model=self.model)
        self.attack_config = attack_config or {}
        self.attack_params: Dict[str, Any] = {}
        
        # Initialize defenses
        self.defenses = self._initialize_defenses()

        # Initialize dataset
        self.train_loader, self.test_loader = self._initialize_data_loaders()

        # Calculate original accuracy
        self.original_accuracy = self.calculate_accuracy(self.test_loader)
        self.current_accuracy = self.original_accuracy
        self.attack_accuracy = self.original_accuracy

        # Define action and observation space
        self.action_space = spaces.Discrete(len(self.defenses))
        # Observation: [current_accuracy, original_accuracy, attack_accuracy, impact_of_attack]
        # Extend this with additional features if needed
        self.observation_space = spaces.Box(
            low=0.0,
            high=1.0,
            shape=(4,),
            dtype=np.float32
        )

        # Visualization components
        self.screen = None
        self.clock = None

    def _initialize_defenses(self) -> Dict[str, Any]:
        """
        Initialize available defense mechanisms.

        Returns:
            Dict[str, Any]: Dictionary of defense mechanisms.
        """
        defenses = {
            'GaussianNoise': GaussianNoiseDefense(sigma=0.02),
            'LabelSmoothing': LabelSmoothingDefense(alpha=0.1, num_classes=1000),
            'Mixup': MixupDefense(alpha=0.4),
            'FeatureSqueezing': FeatureSqueezingDefense(bit_depth=4),
            'SpatialSmoothing': SpatialSmoothingDefense(kernel_size=5),
            'Cutout': CutoutDefense(mask_size=50),
            'JpegCompression': None  # Initialized conditionally
        }

        try:
            from PIL import Image  # noqa: F401
            defenses['JpegCompression'] = JpegCompressionDefense(quality=50)
        except ImportError:
            logger.error("PIL is required for JpegCompressionDefense. Please install it using `pip install Pillow`.")

        logger.info(f"Initialized defenses: {list(defenses.keys())}")
        return defenses

    def _initialize_data_loaders(self) -> Tuple[Any, Any]:
        """
        Initialize data loaders for training and testing.

        Returns:
            Tuple[Any, Any]: Training and testing data loaders.
        """
        try:
            data_transforms = transforms.Compose([
                transforms.Lambda(lambda img: img.convert("RGB")),
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                )
            ])

            train_loader, test_loader = datapipeline(
                config=self.config,
                transforms_fn=data_transforms,
                target_column="image",
                label_column="label",
            )
            logger.info("Data loaders created successfully.")
            return train_loader, test_loader
        except Exception as e:
            logger.error(f"Failed to create data loaders: {e}")
            raise e

    def calculate_accuracy(self, data_loader: Any) -> float:
        """
        Calculate the model's accuracy on a given data loader.

        Args:
            data_loader (Any): Data loader for evaluation.

        Returns:
            float: Accuracy score.
        """
        try:
            correct = 0
            total = 0
            with torch.no_grad():
                for batch in data_loader:
                    images, labels = batch["images"].to(self.model_pipeline.device), batch["labels"].to(self.model_pipeline.device)
                    outputs = self.model(images).logits
                    _, predicted = torch.max(outputs, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            accuracy = correct / total if total > 0 else 0.0
            logger.info(f"Calculated accuracy: {accuracy:.4f}")
            return accuracy
        except Exception as e:
            logger.error(f"Failed to calculate accuracy: {e}")
            return 0.0

    def reset(self) -> np.ndarray:
        """
        Reset the environment to the initial state.

        Returns:
            np.ndarray: Initial observation.
        """
        self.current_accuracy = self.original_accuracy
        self.attack_accuracy = self.original_accuracy
        self.attack_params = {}
        logger.info("Environment reset.")
        return self._get_observation()

    def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]:
        """
        Execute an action (defense strategy) in the environment.

        Args:
            action (int): Defense strategy index to apply.

        Returns:
            Tuple[np.ndarray, float, bool, Dict[str, Any]]:
                Observation after action,
                Reward,
                Done flag,
                Additional info.
        """
        try:
            defense_name = list(self.defenses.keys())[action]
            logger.info(f"Applying defense: {defense_name}")
            defense = self.defenses.get(defense_name)
            if defense is None:
                logger.warning(f"Defense '{defense_name}' not found or not initialized. Penalizing.")
                return self._get_observation(), -1.0, False, {"error": "Invalid defense"}

            # Apply attack
            attacked_images, attacked_labels = self._apply_attack()

            # Apply defense
            defended_images, defended_labels = self._apply_defense(defense, attacked_images, attacked_labels)

            # Evaluate defended model
            defended_accuracy = self.evaluate_defended_model(defended_images, defended_labels)

            # Update accuracies
            self.attack_accuracy = defended_accuracy
            impact = self.original_accuracy - self.attack_accuracy
            self.current_accuracy = self.attack_accuracy

            # Calculate reward
            reward = self.calculate_reward(impact)

            # Define done condition (e.g., when a certain accuracy threshold is met)
            done = self._is_done(impact)

            logger.info(f"Defense '{defense_name}' applied. Defended accuracy: {defended_accuracy:.4f}, Reward: {reward:.4f}")

            observation = self._get_observation()
            info = {"defense_applied": defense_name, "impact": impact}

            return observation, reward, done, info

        except Exception as e:
            logger.error(f"Error during step: {e}")
            return self._get_observation(), -1.0, False, {"error": str(e)}

    def _apply_attack(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Apply a set of adversarial attacks to the dataset.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Attacked images and labels.
        """
        try:
            # Select a random attack type
            attack_name = np.random.choice(list(self.adversary.get_supported_attacks()))
            logger.info(f"Applying attack: {attack_name}")

            # Configure attack parameters if any
            config_attack: Optional[AttackConfig] = self.attack_config.get(attack_name, {})

            # Iterate over test data to apply attack
            attacked_images_list = []
            attacked_labels_list = []

            for batch in self.test_loader:
                images, labels = batch["images"].to(self.model_pipeline.device), batch["labels"].to(self.model_pipeline.device)
                adv_images = self.adversary.attack(attack_name, images, labels, config_attack)
                attacked_images_list.append(adv_images)
                attacked_labels_list.append(labels)

            # Concatenate all batches
            attacked_images = torch.cat(attacked_images_list, dim=0)
            attacked_labels = torch.cat(attacked_labels_list, dim=0)

            self.attack_params = {'type': attack_name}

            logger.info(f"Attack '{attack_name}' applied successfully.")
            return attacked_images, attacked_labels

        except Exception as e:
            logger.error(f"Failed to apply attack: {e}")
            raise e

    def _apply_defense(
        self,
        defense: Any,
        images: torch.Tensor,
        labels: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Apply the selected defense to the images and labels.

        Args:
            defense (Any): Defense mechanism instance.
            images (torch.Tensor): Batch of images.
            labels (torch.Tensor): Batch of labels.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Defended images and labels.
        """
        try:
            if isinstance(defense, GaussianNoiseDefense):
                defended_images = defense.apply(images)
                return defended_images, labels
            elif isinstance(defense, LabelSmoothingDefense):
                defended_labels = defense.apply(labels)
                return images, defended_labels
            elif isinstance(defense, MixupDefense):
                defended_images, defended_labels = defense.apply(images, labels)
                return defended_images, defended_labels
            elif isinstance(defense, FeatureSqueezingDefense):
                defended_images = defense.apply(images)
                return defended_images, labels
            elif isinstance(defense, SpatialSmoothingDefense):
                defended_images = defense.apply(images)
                return defended_images, labels
            elif isinstance(defense, CutoutDefense):
                defended_images = defense.apply(images)
                return defended_images, labels
            elif isinstance(defense, JpegCompressionDefense):
                defended_images = defense.apply(images)
                return defended_images, labels
            else:
                logger.warning(f"Defense type '{type(defense)}' not recognized. Skipping defense.")
                return images, labels
        except Exception as e:
            logger.error(f"Failed to apply defense '{type(defense)}': {e}")
            return images, labels

    def evaluate_defended_model(self, images: torch.Tensor, labels: torch.Tensor) -> float:
        """
        Evaluate the model's accuracy on the defended data.

        Args:
            images (torch.Tensor): Defended images.
            labels (torch.Tensor): Defended labels.

        Returns:
            float: Defended accuracy.
        """
        try:
            correct = 0
            total = 0
            with torch.no_grad():
                outputs = self.model(images).logits
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
            accuracy = correct / total if total > 0 else 0.0
            logger.info(f"Defended model accuracy: {accuracy:.4f}")
            return accuracy
        except Exception as e:
            logger.error(f"Failed to evaluate defended model: {e}")
            return 0.0

    def _get_observation(self) -> np.ndarray:
        """
        Get the current observation from the environment.

        Returns:
            np.ndarray: Current state observation.
        """
        observation = np.array([
            self.current_accuracy,
            self.original_accuracy,
            self.attack_accuracy,
            self.original_accuracy - self.attack_accuracy
        ], dtype=np.float32)
        return observation

    def calculate_reward(self, impact: float) -> float:
        """
        Calculate the reward based on the impact of the defense.

        Args:
            impact (float): Difference between original and attack accuracies.

        Returns:
            float: Calculated reward.
        """
        try:
            # Reward mechanism:
            # Higher impact (greater restoration of accuracy) yields higher rewards
            # Optionally, penalize for using complex or resource-intensive defenses
            reward = impact
            logger.debug(f"Calculated reward: {reward}")
            return reward
        except Exception as e:
            logger.error(f"Failed to calculate reward: {e}")
            return -1.0

    def _is_done(self, impact: float) -> bool:
        """
        Determine if the episode is done based on the impact.

        Args:
            impact (float): Difference between original and attack accuracies.

        Returns:
            bool: Whether the episode is done.
        """
        # Define your termination condition
        # Example: Stop if accuracy restoration exceeds a threshold or max steps reached
        accuracy_threshold = 0.95  # Example threshold
        return self.current_accuracy >= accuracy_threshold

    def render(self, mode: str = 'human') -> None:
        """
        Render the environment visually.

        Args:
            mode (str, optional): The mode to render with. Defaults to 'human'.
        """
        if mode == 'human':
            if self.screen is None:
                import pygame
                pygame.init()
                self.screen = pygame.display.set_mode((400, 300))
                self.clock = pygame.time.Clock()

            self.screen.fill((255, 255, 255))  # White background

            # Display current metrics
            font = pygame.font.SysFont(None, 24)
            metrics_text = f"Original Acc: {self.original_accuracy:.2f} | Current Acc: {self.current_accuracy:.2f} | Attack Acc: {self.attack_accuracy:.2f}"
            text = font.render(metrics_text, True, (0, 0, 0))
            self.screen.blit(text, (20, 150))

            pygame.display.flip()
            self.clock.tick(10)

    def close(self) -> None:
        """
        Clean up resources.
        """
        if self.screen is not None:
            import pygame
            pygame.quit()
            self.screen = None
            self.clock = None