class KernelShapExplainer(FeatureAttributionExplainer):
    """Kernel SHAP explainer for computing Shapley values using LIME framework.

    This explainer computes attributions using Kernel SHAP, which uses the LIME framework
    with specific weighting and sampling strategies to efficiently compute Shapley values.
    Kernel SHAP provides theoretically grounded explanations that satisfy Shapley value
    axioms (efficiency, symmetry, dummy, additivity) while being more computationally
    efficient than direct Shapley value computation. Supports both single-target and
    multi-target modes for both single-target and multi-target scenarios.

    Kernel SHAP is particularly effective for tabular data and provides globally
    consistent explanations across different inputs.

    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 1.
        n_samples: Number of coalition samples for Shapley value estimation.
            Defaults to 100.
        weight_attributions: Whether to weight attributions by feature group sizes
            when using feature masks. Defaults to True.

    Examples:
        Single-target usage without mask (feature-level):
        >>> import torch
        >>> from torchxai.data_types import SingleTargetAcrossBatch
        >>>
        >>> model = torch.nn.Linear(10, 2)
        >>> explainer = KernelShapExplainer(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 = KernelShapExplainer(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",
        "_n_samples",
        "_weight_attributions",
    ]

    def __init__(
        self,
        model: Module,
        multi_target: bool = False,
        internal_batch_size: int = 1,
        n_samples: int = 100,
        weight_attributions: bool = True,
    ) -> None:
        """Initialize the KernelShapExplainer.

        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.
            n_samples: Number of coalition samples for Shapley estimation. Defaults to 100.
            weight_attributions: Whether to weight attributions by feature groups. Defaults to True.
        """
        self._n_samples = n_samples
        self._weight_attributions = weight_attributions

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

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

        Returns:
            KernelShap attribution function for single targets.
        """
        expl_func = partial(
            KernelShap(self._model).attribute,
            n_samples=self._n_samples,
            perturbations_per_eval=self._internal_batch_size,
        )
        return self._expl_fn_with_post_process(expl_func)

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

        Returns:
            MultiTargetKernelShap attribution function for multiple targets.
        """
        expl_func = partial(
            MultiTargetKernelShap(self._model).attribute,
            n_samples=self._n_samples,
            perturbations_per_eval=self._internal_batch_size,
        )
        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,
        frozen_features: list[torch.Tensor] | None = None,
    ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
        """Compute Kernel SHAP 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 coalition sampling (typically zeros).
                If None, uses zero baselines matching input shape.
            feature_mask: Masks representing feature groups for aggregation. Features
                with the same mask value are treated as a single coalition member.
            additional_forward_args: Additional arguments for model forward pass.
            frozen_features: List of feature indices to keep unchanged during perturbation.
                Useful for special tokens like CLS, SEP in NLP models.

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

        Note:
            Kernel SHAP uses coalition sampling with SHAP-specific weighting to estimate
            Shapley values. The number of samples significantly affects accuracy and
            computation time. More samples provide better Shapley value approximations.

        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,
            frozen_features=frozen_features,
        )