# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
# See the LICENSE file in the root of this repository for details.
# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial

import warnings

import plotly.express as px
from datasets import Dataset

from validmind import tags, tasks
from validmind.errors import MissingDependencyError

from .utils import get_ragas_config, get_renamed_columns

try:
    from ragas import evaluate
    from ragas.metrics import answer_relevancy
except ImportError as e:
    raise MissingDependencyError(
        "Missing required package `ragas` for AnswerRelevance. "
        "Please run `pip install validmind[llm]` to use LLM tests",
        required_dependencies=["ragas"],
        extra="llm",
    ) from e


@tags("ragas", "llm", "rag_performance")
@tasks("text_qa", "text_generation", "text_summarization")
def AnswerRelevance(
    dataset,
    question_column="question",
    contexts_column="contexts",
    answer_column="answer",
):
    """
    Assesses how pertinent the generated answer is to the given prompt.

    The evaluation metric, Answer Relevancy, focuses on assessing how pertinent the
    generated answer is to the given prompt. A lower score is assigned to answers that
    are incomplete or contain redundant information and higher scores indicate better
    relevancy. This metric is computed using the `question`, the `contexts` and the
    `answer`.

    The Answer Relevancy is defined as the mean cosine similartiy of the original
    `question` to a number of artifical questions, which are generated (reverse-engineered)
    based on the `answer`:

    $$
    \\text{answer relevancy} = \\frac{1}{N} \\sum_{i=1}^{N} cos(E_{g_i}, E_o)
    $$
    $$
    \\text{answer relevancy} = \\frac{1}{N} \\sum_{i=1}^{N} \\frac{E_{g_i} \\cdot E_o}{\\|E_{g_i}\\|\\|E_o\\|}
    $$

    Where:
    - $E_{g_i}$ is the embedding of the generated question $i$.
    - $E_o$ is the embedding of the original question.
    - $N$ is the number of generated questions - 3 by default.

    **Note**: *This is a reference-free metric, meaning that it does not require a
    `ground_truth` answer to compare against. A similar metric that does evaluate the
    correctness of a generated answser with respect to a `ground_truth` answer is
    `validmind.model_validation.ragas.AnswerCorrectness`.*

    ### Configuring Columns

    This metric requires the following columns in your dataset:

    - `question` (str): The text query that was input into the model.
    - `contexts` (List[str]): Any contextual information retrieved by the model before
    generating an answer.
    - `answer` (str): The response generated by the model.

    If the above data is not in the appropriate column, you can specify different column
    names for these fields using the parameters `question_column`, `answer_column`, and
    `contexts_column`.

    For example, if your dataset has this data stored in different columns, you can
    pass the following parameters:
    ```python
    params = {
        "question_column": "input_text",
        "answer_column": "output_text",
        "contexts_column": "context_info"
    }
    ```

    If answer and contexts are stored as a dictionary in another column, specify the
    column and key like this:
    ```python
    pred_col = dataset.prediction_column(model)
    params = {
        "answer_column": f"{pred_col}.generated_answer",
        "contexts_column": f"{pred_col}.contexts",
    }
    ```

    For more complex data structures, you can use a function to extract the answers:
    ```python
    pred_col = dataset.prediction_column(model)
    params = {
        "answer_column": lambda row: "\\n\\n".join(row[pred_col]["messages"]),
        "contexts_column": lambda row: [row[pred_col]["context_message"]],
    }
    ```
    """
    warnings.filterwarnings(
        "ignore",
        category=FutureWarning,
        message="promote has been superseded by promote_options='default'.",
    )

    required_columns = {
        "question": question_column,
        "answer": answer_column,
        "contexts": contexts_column,
    }

    df = get_renamed_columns(dataset._df, required_columns)

    result_df = evaluate(
        Dataset.from_pandas(df), metrics=[answer_relevancy], **get_ragas_config()
    ).to_pandas()

    fig_histogram = px.histogram(x=result_df["answer_relevancy"].to_list(), nbins=10)
    fig_box = px.box(x=result_df["answer_relevancy"].to_list())

    return (
        {
            # "Scores (will not be uploaded to UI)": result_df[
            #     ["question", "contexts", "answer", "answer_relevancy"]
            # ],
            "Aggregate Scores": [
                {
                    "Mean Score": result_df["answer_relevancy"].mean(),
                    "Median Score": result_df["answer_relevancy"].median(),
                    "Max Score": result_df["answer_relevancy"].max(),
                    "Min Score": result_df["answer_relevancy"].min(),
                    "Standard Deviation": result_df["answer_relevancy"].std(),
                    "Count": result_df.shape[0],
                }
            ],
        },
        fig_histogram,
        fig_box,
    )
