"""Explainers are defined here."""

import numpy as np
import polars as pl
import shap
from scipy.spatial import cKDTree


class ShapExplainer:
    """Defines a SHAP explainer for reference point based methods."""

    def __init__(self, problem_data: pl.DataFrame, input_symbols: list[str], output_symbols: list[str]):
        """Initialize the explainer.

        Initializes the explainer with given data, and input and output symbols.
        The data should contain the columns listed in the input and output symbols.
        This data is then used to simulate the inputs and outputs of an (interactive)
        multiobjective optimization method, which is used to explain the relation of its
        inputs and outputs using SHAP values.

        Note:
            The `data` can be generated by for a reference point based based by, e.g.,
            randomly sampling the input space and then evaluating the methods with the
            sampled inputs to generate outputs.

        Args:
            problem_data (pl.DataFrame): the data to simulate the input and
                outputs of a multiobjective optimization method.
            input_symbols (list[str]): the input symbols present in `data`.
                These symbols represent the inputs to the method.
            output_symbols (list[str]): the output symbols present in `data`.
                These symbols represent the outputs of the method.
        """
        self.data = problem_data
        self.input_symbols = input_symbols
        self.output_symbols = output_symbols
        self.input_array = self.data[self.input_symbols].to_numpy()
        self.output_array = self.data[self.output_symbols].to_numpy()
        self.to_output_tree = cKDTree(self.input_array)
        self.explainer = None

    def setup(self, background_data: pl.DataFrame):
        """Setup the explainer.

        Setups the SHAP explainer with  the given background data. The
        background data should have the columns `self.input_symbols`. The
        background data is used as the background (or missing data) when
        computing SHAP values.  The mean (or expected values) of the background
        data's output (`self.output_symbols`) will determine the baseline of the
        SHAP values.

        Note:
            To generate a dataset with meaningful expected values, e.g., in case
            the SHAP values are better understood by relating them to a specific baseline,
            see `desdeo.explanations.generate_biased_mean_data`.

        Args:
            background_data (pl.DataFrame): the background data.
        """
        self.explainer = shap.Explainer(
            self.evaluate,
            masker=background_data[self.input_symbols].to_numpy(),
        )

    def evaluate(self, evaluate_array: np.ndarray) -> np.ndarray:
        """Evaluates the multiobjective optimization method represented by the data.

        Note:
            Evaluation happens by finding the closest matching input array in the
            `self.input_array` and then using that value's corresponding output
            as the evaluation result. Closest means lowest Euclidean distance.

        Args:
            evaluate_array (np.ndarray): the inputs to the method represented by the data.
                Can be either a single input, or an array of multiple inputs. Used mainly by
                `self.explain_input`.

        Returns:
            np.ndarray: the evaluated output(s) corresponding to the input data.
        """
        _, indices = self.to_output_tree.query(evaluate_array)

        return self.output_array[indices]

    def explain_input(self, to_be_explained: pl.DataFrame) -> dict:
        """Explain an input and produces SHAP values.

        Args:
            to_be_explained (pl.DataFrame): the input to be explained. The
                dataframe must have the columns defined in `self.input_symbols`.

        Returns:
            dict: the key 'shaps' corresponds to the computed SHAP values for
                the input, the key 'base_values' is the baseline the SHAP values
                were computed against, and the key 'data' is the input the SHAP
                values were computed for.
        """
        _to_be_explained = to_be_explained[self.input_symbols].to_numpy()

        return self.explainer(_to_be_explained)
