class GradientShapExplainer(FeatureAttributionExplainer):
    """GradientShap explainer for computing noise-based Shapley value approximations.

    This explainer computes attributions using GradientShap, which combines ideas from
    Integrated Gradients and SHAP. It uses a distribution of baselines rather than a
    single baseline and adds noise to create more robust attribution estimates. The method
    approximates Shapley values through gradient-based computations. Supports both
    single-target and multi-target modes for both single-target and multi-target scenarios.

    GradientShap provides more robust attributions by using baseline distributions
    and noise, making it less sensitive to specific baseline choices.

    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. Defaults to 64.
        grad_batch_size: Batch size for gradient computations. Defaults to 64.
        n_samples: Number of random samples used to approximate Shapley values.
            Defaults to 25.
        return_convergence_delta: Whether to return convergence delta for
            completeness check. Defaults to False.

    Examples:
        Single-target usage:
        >>> import torch
        >>> from torchxai.data_types import SingleTargetAcrossBatch
        >>>
        >>> model = torch.nn.Linear(10, 2)
        >>> explainer = GradientShapExplainer(model)
        >>> inputs        = torch.randn(1, 10)
        >>> baselines_dist = torch.zeros(1, 10).expand(5, -1)   # 5-sample reference distribution
        >>> attributions = explainer.explain(
        ...     inputs=inputs,
        ...     baselines=baselines_dist,
        ...     target=SingleTargetAcrossBatch(index=0),
        ... )
        >>> attributions.shape   # (1, 10)

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

    __repr_attrs__ = [
        "_multi_target",
        "_internal_batch_size",
        "_grad_batch_size",
        "_n_samples",
        "_return_convergence_delta",
        "_stdevs",
    ]

    def __init__(
        self,
        model: Module,
        multi_target: bool = False,
        internal_batch_size: int = 1,
        grad_batch_size: int = 64,
        n_samples: int = 25,
        stddevs: float = 0.0,
        return_convergence_delta: bool = False,
    ) -> None:
        """Initialize the GradientShapExplainer.

        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 1.
            grad_batch_size: Batch size for gradient computations. Defaults to 64.
            n_samples: Number of random samples for Shapley value approximation.
                Defaults to 25.
            return_convergence_delta: Whether to return convergence delta for
                completeness check. Defaults to False.
        """
        self._n_samples = n_samples
        self._return_convergence_delta = return_convergence_delta
        self._stdevs = stddevs

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

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

        Returns:
            Custom GradientShap attribution function for single targets.
        """
        return partial(
            GradientShap_(self._model).attribute,
            n_samples=self._n_samples,
            n_samples_batch_size=self._internal_batch_size,
            return_convergence_delta=self._return_convergence_delta,
            stdevs=self._stdevs,
        )

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

        Returns:
            MultiTargetGradientShap attribution function for multiple targets.
        """
        return partial(
            MultiTargetGradientShap(
                self._model, grad_batch_size=self._grad_batch_size
            ).attribute,
            n_samples=self._n_samples,
            n_samples_batch_size=self._internal_batch_size,
            return_convergence_delta=self._return_convergence_delta,
            stdevs=self._stdevs,
        )

    def explain(
        self,
        inputs: TensorOrTupleOfTensorsGeneric,
        target: ExplanationTargetType | list[ExplanationTargetType],
        baselines: TensorOrTupleOfTensorsGeneric | None = None,
        additional_forward_args: tuple[Any, ...] | None = None,
    ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
        """Compute GradientShap 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 distribution for GradientShap. Must be provided as
                a tensor distribution or callable that generates baseline samples.
            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:
            GradientShap requires a distribution of baselines rather than a single baseline.
            The number of samples and noise parameters are controlled by initialization settings.

        Examples:
            >>> # With baseline distribution
            >>> baseline_dist = torch.randn(50, 10)  # 50 baseline samples
            >>> attributions = explainer.explain(
            ...     inputs=OrderedDict({"input": torch.randn(2, 10)}),
            ...     target=torch.tensor([0, 1]),
            ...     baselines=OrderedDict({"input": baseline_dist}),
            ... )
        """
        return self._default_explain(
            inputs=inputs,
            target=target,
            baselines=baselines,
            additional_forward_args=additional_forward_args,
        )