"""Base perturbator classes for applying masks to create perturbed inputs.

This module provides the foundational infrastructure for perturbators in the MUPPET XAI
framework. Perturbators implement the second step of the four-step perturbation-based
explanation process: generate masks → **apply perturbations** → calculate attributions
→ aggregate results.

Perturbators take the original input data and binary masks from explorers, then create
perturbed versions of the input by modifying specific regions according to the masks.
The perturbation strategy depends on the data modality and explanation method requirements.

The module contains:
    Perturbator: Abstract base class defining the perturbation interface with automatic
        batch processing and memory management capabilities.
    TrainablePerturbator: Extended base class for perturbators that use trainable
        generators for more sophisticated perturbation strategies.

Mask Convention:
    All premise types follow the MUPPET binary mask convention:
    - 0: Preserve the original input value (no perturbation)
    - 1: Perturb the input value (apply perturbation strategy)

Note:
    Perturbators work closely with explorers (mask generation) and are consumed by
    attributors (feature importance calculation). The choice of perturbation strategy
    significantly impacts explanation quality.
"""

#
# Created on Wed Jun 14 2023
#
# Copyright (c) 2023 Ismail Bachchar @Euranova
#

import functools
from abc import ABC, abstractmethod

import torch
from torch.utils.data import DataLoader

from muppet import logger
from muppet.components.perturbator.generator.base import TrainableGenerator


class Perturbator(ABC):
    """Abstract base class for perturbation strategies in XAI methods.

    Perturbators apply masks to input data to create perturbed versions
    for explanation purposes. They define how the original input should
    be modified based on the perturbation masks generated by explorers.

    Perturbators implement the second step of the four-step perturbation-based
    explanation process: generate masks → **apply perturbations** → calculate
    attributions → aggregate results.

    The perturbators automatically handle:
        - Batch processing with configurable memory limits
        - Out-of-memory error recovery with smaller batch sizes
        - Integration with trainable generators when needed

    """

    def __init__(self, max_batch_size: int) -> None:
        """Initialize the Perturbator with batch processing configuration.

        Args:
            max_batch_size (int): Maximum batch size for processing perturbations.
        """
        self.device = None
        self.max_batch_size = max_batch_size
        super().__init__()

    def reinitialize(self):
        """Reset the perturbator to its initial state.

        Clears any internal state or cached data that may affect
        subsequent perturbation operations.
        """
        pass

    @abstractmethod
    def perturbate(
        self,
        x: torch.Tensor,
        masks: torch.Tensor,
    ) -> torch.Tensor:
        """Apply perturbations to input data using provided masks.

        Mask Convention:
            All perturbators follow the MUPPET perturbation convention:
            - 0: Preserve the original input value (no perturbation)
            - 1: Perturb the input value (apply perturbation strategy)

        Args:
            x (torch.Tensor): The original input example. Shape (b, *x.shape[1:]).
            masks (torch.Tensor): Binary masks for perturbation. Shape (N, *mask_shape).

        Returns:
            torch.Tensor: Perturbed examples. Shape (N, *x.shape).
        """
        raise NotImplementedError

    @staticmethod
    def batch_perturbation():
        """A method that acts as a decorator factory.
        It has access to the instance 'self'.
        """

        def decorator(func):
            """Decorator to automatically handle torch.OutOfMemoryError by
            retrying the function with smaller batches.

            This decorator assumes:
            1. The function is a method of a class.
            2. The second argument (index 1) is the tensor to be batched
            along its first dimension (e.g., 'masks').
            3. The function returns a torch.Tensor that can be concatenated.
            """

            @functools.wraps(func)
            def wrapper(instance_self, x, masks, *args, **kwargs):
                # the tensor to be batched ('masks') from args.
                # This assumes the function signature is like: method(self, x, masks, ...)

                # The initial batch size is the total number of items.
                tensor_to_batch = masks
                n = tensor_to_batch.size(0)
                batch_size = n

                if batch_size >= instance_self.max_batch_size:
                    new_batch_size = instance_self.max_batch_size
                    logger.debug(
                        f"Reducing batch size from {batch_size} to {new_batch_size} preventively"
                    )
                    batch_size = new_batch_size

                results = []
                logger.debug(
                    f"Attempting to run '{func.__name__}' with batch size {batch_size}..."
                )
                # Process the input tensor in chunks of the current batch_size
                for i in range(0, n, batch_size):
                    chunk = tensor_to_batch[i : i + batch_size]

                    result_chunk = func(
                        instance_self, x, chunk, *args, **kwargs
                    )
                    results.append(result_chunk)

                # If all chunks were processed successfully, concatenate and return.
                logger.debug(
                    f"Successfully processed '{func.__name__}' with batch size {batch_size}."
                )
                return torch.cat(results, dim=0)

            return wrapper

        return decorator

    @batch_perturbation()
    def __call__(self, x: torch.Tensor, masks: torch.Tensor):
        """Apply perturbation to input data using the provided masks.

        This method serves as the callable interface for the perturbator,
        applying the perturbation method to the input data using the specified
        masks and ensuring the result is on the correct device.

        Mask Convention:
            All perturbators follow the MUPPET binary mask convention:
            - 0: Preserve the original input value (no perturbation)
            - 1: Perturb the input value (apply perturbation strategy)

        Args:
            x (torch.Tensor): The input data tensor to be perturbed.
            masks (torch.Tensor): The binary masks to apply during perturbation.

        Returns:
            torch.Tensor: The perturbed data tensor on the correct device.
        """
        perturbations = self.perturbate(x=x, masks=masks)
        perturbations = perturbations.to(self.device)
        return perturbations


class TrainablePerturbator(Perturbator):
    """Perturbator with trainable generator capabilities.

    Extends the basic Perturbator to include a trainable generator
    component that can learn perturbation strategies from training data.
    These trainable perturbators are designed for perturbation strategies
    that create plausible perturbations.

    It automatically handles generator training if needed and integrates
    seamlessly with the MUPPET explanation pipeline.

    Technical Features:
        - Automatic generator training with validation splits and early stopping
        - Seamless integration with various trainable generators
    """

    def __init__(
        self,
        generator: TrainableGenerator,
        train_loader: DataLoader,
        max_batch_size: int = 100,
    ) -> None:
        """Initialize the TrainablePerturbator with a generator and training data.

        Args:
            generator (TrainableGenerator): The generator for creating perturbations.
            train_loader (DataLoader): Training data for generator training.
            max_batch_size (int): Maximum batch size for processing. Defaults to 100.
        """
        self.generator = generator
        if not self.generator.is_trained:
            assert train_loader is not None, (
                "If the generator training is required, you must provide the training data!"
            )
            train_losses, test_losses = self.generator.train_generator(
                train_loader=train_loader,
            )
            logger.info(
                "Finished training generator on provided data! Accessible through: 'perturbator.generator'."
            )
            logger.debug(
                "Training loss trend: " + str([i for i in train_losses])
            )
            logger.debug("Testing loss trend: " + str([i for i in test_losses]))

            # set generator state to trained
            self.generator.is_trained = True

        super().__init__(max_batch_size=max_batch_size)
