class FeatureAblationExplainer(FeatureAttributionExplainer):
    """Feature Ablation explainer for computing systematic feature removal attributions.

    This explainer computes attributions using Feature Ablation, which systematically
    removes features or feature groups and measures the resulting change in model output.
    This direct approach provides intuitive explanations by showing exactly how much
    each feature contributes to the prediction. The method supports feature grouping
    through masks, allowing for hierarchical ablation studies. Supports both single-target
    and multi-target modes for both single-target and multi-target scenarios.

    Feature Ablation provides the most direct measure of feature importance through
    systematic removal and impact measurement.

    Args:
        model: The PyTorch model whose output is to be explained.
        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 64.
        weight_attributions: Whether to weight attributions by feature group sizes
            when using feature masks. Defaults to False.

    Examples:
        Single-target usage without mask (feature-level):
        >>> import torch
        >>> from torchxai.data_types import SingleTargetAcrossBatch
        >>>
        >>> model = torch.nn.Linear(10, 2)
        >>> explainer = FeatureAblationExplainer(model)
        >>> attributions = explainer.explain(
        ...     inputs=torch.randn(1, 10),
        ...     target=SingleTargetAcrossBatch(index=0),
        ... )
        >>> attributions.shape   # (1, 10)

        With a feature mask (group-level attribution):
        >>> feature_mask = torch.tensor([[0, 0, 1, 1, 2, 2, 2, 3, 3, 4]])
        >>> attributions_grouped = explainer.explain(
        ...     inputs=torch.randn(1, 10),
        ...     feature_mask=feature_mask,
        ...     target=SingleTargetAcrossBatch(index=0),
        ... )
        >>> attributions_grouped.shape   # (1, 10)

        Multi-target usage:
        >>> explainer_mt = FeatureAblationExplainer(model, multi_target=True)
        >>> mt_attributions = explainer_mt.explain(
        ...     inputs=torch.randn(1, 10),
        ...     target=[SingleTargetAcrossBatch(index=0), SingleTargetAcrossBatch(index=1)],
        ... )
        >>> len(mt_attributions), mt_attributions[0].shape   # 2, (1, 10)
    """

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

    def __init__(
        self,
        model: torch.nn.Module,
        multi_target: bool = False,
        internal_batch_size: int = 64,
        weight_attributions: bool = False,
        show_progress: bool = False,
    ) -> None:
        """Initialize the FeatureAblationExplainer.

        Args:
            model: The model whose output is to be explained.
            multi_target: Whether to use multi-target mode. Defaults to False.
            internal_batch_size: Batch size for internal computations. Defaults to 64.
            weight_attributions: Whether to weight attributions by feature groups. Defaults to False.
        """
        self._weight_attributions = weight_attributions
        self._show_progress = show_progress

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

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

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

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

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

    def _expl_fn_with_post_process(self, expl_func: Callable) -> Callable:
        def wrapped(
            inputs: TensorOrTupleOfTensorsGeneric,
            baselines: BaselineType = None,
            target: TargetType = None,
            additional_forward_args: Any = None,
            feature_mask: None | Tensor | tuple[Tensor, ...] = None,
            **kwargs,
        ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
            feature_mask = (
                _expand_feature_mask_to_target(feature_mask, inputs)
                if feature_mask is not None
                else None
            )
            additional_forward_args = _format_additional_forward_args(
                additional_forward_args
            )

            attributions = expl_func(
                inputs=inputs,
                baselines=baselines,
                target=target,
                additional_forward_args=additional_forward_args,
                feature_mask=feature_mask,
                **kwargs,
            )

            # Apply feature weighting if requested
            if self._weight_attributions and feature_mask is not None:
                if self._multi_target:
                    attributions = [
                        _weight_attributions(attribution, feature_mask)
                        for attribution in attributions
                    ]
                else:
                    attributions = _weight_attributions(attributions, feature_mask)
            return attributions

        return wrapped

    def explain(
        self,
        inputs: TensorOrTupleOfTensorsGeneric,
        target: ExplanationTargetType | list[ExplanationTargetType],
        baselines: TensorOrTupleOfTensorsGeneric | None = None,
        feature_mask: TensorOrTupleOfTensorsGeneric | None = None,
        additional_forward_args: tuple[Any, ...] | None = None,
    ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
        """Compute Feature Ablation 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 ablation (typically zeros). If None,
                uses zero baselines matching input shape.
            feature_mask: Masks representing feature groups for ablation. Features
                with the same mask value are ablated together as a group.
            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:
            Feature Ablation directly measures feature importance through systematic removal.
            The internal_batch_size parameter controls how many features are ablated
            simultaneously for computational efficiency.

        Examples:
            >>> # For tabular data with feature grouping
            >>> feature_mask = torch.tensor([[0, 0, 1, 1, 2, 2, 2, 3, 3, 4]])
            >>> attributions = explainer.explain(
            ...     inputs=OrderedDict({"features": torch.randn(1, 10)}),
            ...     target=torch.tensor([1]),
            ...     baselines=OrderedDict({"features": torch.zeros(1, 10)}),
            ...     feature_mask=feature_mask,
            ... )
        """
        return self._default_explain(
            inputs=inputs,
            target=target,
            baselines=baselines,
            feature_mask=feature_mask,
            additional_forward_args=additional_forward_args,
        )