class OcclusionExplainer(FeatureAttributionExplainer):
    """Occlusion explainer for computing sliding-window perturbation attributions.

    This explainer computes attributions using the Occlusion method, which systematically
    replaces rectangular regions of the input with baseline values (typically zeros)
    and measures the resulting change in model output. This approach is particularly
    effective for image data where spatial regions can be meaningfully occluded.
    Supports both single-target and multi-target modes for both single-target and multi-target scenarios.

    The Occlusion method provides intuitive attributions by directly measuring
    the importance of spatial regions through systematic perturbation.

    Args:
        model: The PyTorch model whose output is to be explained.
        sliding_window_shapes: Shape of the occlusion window for each input tensor.
            Can be a single tuple for all inputs or tuple of tuples for each input.
        strides: Stride for sliding the occlusion window. Can be int, tuple, or
            tuple of tuples. If None, defaults to sliding_window_shapes.
        multi_target: Whether to use multi-target mode. When True, can compute
            attributions for multiple targets simultaneously. Defaults to False.
        internal_batch_size: Batch size for internal computations (perturbations
            per evaluation). Defaults to 1.

    Examples:
        Single-target usage for image data:
        >>> import torch
        >>> from torchxai.data_types import SingleTargetAcrossBatch
        >>>
        >>> model = torch.nn.Sequential(
        ...     torch.nn.Conv2d(3, 16, 3, padding=1), torch.nn.ReLU(),
        ...     torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(),
        ...     torch.nn.Linear(16, 10),
        ... )
        >>> explainer = OcclusionExplainer(model)
        >>> inputs = torch.randn(1, 3, 32, 32)
        >>> attributions = explainer.explain(
        ...     inputs=inputs,
        ...     sliding_window_shapes=(1, 8, 8),
        ...     target=SingleTargetAcrossBatch(index=0),
        ... )
        >>> attributions.shape   # (1, 3, 32, 32)

        Multi-target usage:
        >>> explainer_mt = OcclusionExplainer(model, multi_target=True)
        >>> mt_attributions = explainer_mt.explain(
        ...     inputs=inputs,
        ...     sliding_window_shapes=(1, 8, 8),
        ...     target=[SingleTargetAcrossBatch(index=0), SingleTargetAcrossBatch(index=1)],
        ... )
        >>> len(mt_attributions), mt_attributions[0].shape   # 2, (1, 3, 32, 32)
    """

    __repr_attrs__ = ["_multi_target", "_internal_batch_size", "_show_progress"]

    def __init__(
        self,
        model: Module,
        multi_target: bool = False,
        internal_batch_size: int = 1,
        show_progress: bool = False,
    ) -> None:
        """Initialize the OcclusionExplainer.

        Args:
            model: The model whose output is to be explained.
            sliding_window_shapes: Shape of the occlusion window for each input tensor.
            strides: Stride for sliding the occlusion window. If None, uses sliding_window_shapes.
            multi_target: Whether to use multi-target mode. Defaults to False.
            internal_batch_size: Batch size for internal computations. Defaults to 1.
        """
        self._show_progress = show_progress

        super().__init__(model, multi_target, internal_batch_size)

    def _init_single_target_explanation_fn(self) -> Callable:
        """Initialize single-target Occlusion attribution function.

        Returns:
            Occlusion attribution function for single targets.
        """
        return partial(
            Occlusion(self._model).attribute,
            perturbations_per_eval=self._internal_batch_size,
            show_progress=self._show_progress,
        )

    def _init_multi_target_explanation_fn(self) -> Callable:
        """Initialize multi-target Occlusion attribution function.

        Returns:
            MultiTargetOcclusion attribution function for multiple targets.
        """
        return partial(
            MultiTargetOcclusion(self._model).attribute,
            perturbations_per_eval=self._internal_batch_size,
            show_progress=self._show_progress,
        )

    def explain(
        self,
        inputs: TensorOrTupleOfTensorsGeneric,
        target: ExplanationTargetType | list[ExplanationTargetType],
        sliding_window_shapes: tuple[int, ...] | tuple[tuple[int, ...], ...],
        strides: None
        | int
        | tuple[int, ...]
        | tuple[int | tuple[int, ...], ...] = None,
        baselines: TensorOrTupleOfTensorsGeneric | None = None,
        additional_forward_args: tuple[Any, ...] | None = None,
    ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
        """Compute Occlusion attributions for the given inputs.

        Args:
            inputs: Input tensor(s) for attribution computation.
            target: An `ExplanationTargetType` (e.g. `SingleTargetAcrossBatch`) for single-target
                mode, or a list of them for multi-target mode.
            baselines: Baseline tensors for occlusion (typically zeros). If None,
                uses zero baselines matching input shape.
            additional_forward_args: Additional arguments for model forward pass.

        Returns:
            Tensor in single-target mode. List of Tensors, one per target, in multi-target mode.

        Note:
            The sliding window and stride parameters are set during initialization.
            Attribution values represent the importance of each spatial region
            as measured by occlusion impact.

        Examples:
            >>> # For image data with 8x8 occlusion windows
            >>> attributions = explainer.explain(
            ...     inputs=OrderedDict({"image": torch.randn(1, 3, 224, 224)}),
            ...     target=torch.tensor([285]),  # ImageNet class
            ...     baselines=OrderedDict({"image": torch.zeros(1, 3, 224, 224)}),
            ... )
        """
        return self._default_explain(
            inputs=inputs,
            target=target,
            sliding_window_shapes=sliding_window_shapes,
            strides=strides,
            baselines=baselines,
            additional_forward_args=additional_forward_args,
        )