class LimeExplainer(FeatureAttributionExplainer):
    """LIME explainer for local interpretable model-agnostic explanations.

    This explainer computes attributions using LIME (Local Interpretable Model-agnostic
    Explanations), which explains individual predictions by learning locally faithful
    linear models around the input. LIME perturbs the input and trains a simple
    interpretable model on the perturbed samples to approximate the model's behavior
    locally. Supports both single-target and multi-target modes for both single-target and multi-target scenarios.

    LIME is particularly useful for understanding complex models by providing locally
    accurate explanations using interpretable linear models.

    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 perturbed samples to generate for LIME. Defaults to 100.
        alpha: Regularization parameter for the LASSO interpretable model. Defaults to 0.01.
        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 = LimeExplainer(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 = LimeExplainer(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",
        "_alpha",
        "_weight_attributions",
        "_show_progress",
    ]

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

        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 perturbed samples for LIME. Defaults to 100.
            alpha: Regularization parameter for LASSO. Defaults to 0.01.
            weight_attributions: Whether to weight attributions by feature groups. Defaults to True.
        """
        self._n_samples = n_samples
        self._alpha = alpha
        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 LIME attribution function.

        Returns:
            LIME attribution function for single targets.
        """
        expl_func = partial(
            Lime(
                self._model,
                interpretable_model=SkLearnLasso(alpha=self._alpha),
                perturb_func=frozen_features_perturb_func,
            ).attribute,
            n_samples=self._n_samples,
            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 LIME attribution function.

        Returns:
            MultiTargetLime attribution function for multiple targets.
        """

        expl_func = partial(
            MultiTargetLime(
                self._model,
                interpretable_model=SkLearnLasso(alpha=self._alpha),
                perturb_func=frozen_features_perturb_func,
            ).attribute,
            n_samples=self._n_samples,
            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,
        frozen_features: list[torch.Tensor] | None = None,
    ) -> TensorOrTupleOfTensorsGeneric | list[TensorOrTupleOfTensorsGeneric]:
        """Compute LIME 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 perturbation (typically zeros). If None,
                uses zero baselines matching input shape.
            feature_mask: Masks representing feature groups for aggregation. If provided,
                features with the same mask value are grouped together.
            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:
            LIME trains a local linear model on perturbed samples. The number of samples
            and regularization strength are controlled by initialization parameters.

        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,
            ... )
        """

        # Get base attributions
        return self._default_explain(
            inputs=inputs,
            target=target,
            baselines=baselines,
            feature_mask=feature_mask,
            additional_forward_args=additional_forward_args,
            frozen_features=frozen_features,
        )