# 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
from typing import Dict, Tuple

import plotly.express as px
import plotly.graph_objects as go
from datasets import Dataset

from validmind import RawData, tags, tasks
from validmind.errors import MissingDependencyError
from validmind.vm_models import VMDataset

from .utils import get_ragas_config, get_renamed_columns

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

    raise e

VALID_FOCUS_VALUES = ["relevant", "irrelevant"]


@tags("ragas", "llm", "rag_performance")
@tasks("text_qa", "text_generation", "text_summarization")
def NoiseSensitivity(
    dataset: VMDataset,
    response_column: str = "response",
    retrieved_contexts_column: str = "retrieved_contexts",
    reference_column: str = "reference",
    focus: str = "relevant",
    user_input_column: str = "user_input",
    judge_llm=None,
    judge_embeddings=None,
) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
    """
    Assesses the sensitivity of a Large Language Model (LLM) to noise in retrieved context by measuring how often it
    generates incorrect responses.

    ### Purpose

    The Noise Sensitivity test aims to measure how sensitive an LLM is to irrelevant or noisy information within the
    contextual data used to generate its responses. A lower noise sensitivity score suggests better model robustness in
    generating accurate answers from given contexts.

    ### Test Mechanism

    This test evaluates the model's answers by comparing the claims made in the generated response against the ground
    truth and the retrieved context. The noise sensitivity score is calculated as:

    $$
    \\text{noise sensitivity} = {|\\text{Number of incorrect claims in answer}| \\over |\\text{Number of total claims in answer}|}
    $$

    The formula computes the fraction of incorrect claims to the total claims in the answer, using a dataset where
    'answer', 'context', and 'ground_truth' columns are specified.

    #### Configuring Columns

    This metric requires the following columns in your dataset:

    - `retrieved_contexts` (List[str]): A list of text contexts which are retrieved to generate
    the answer.
    - `response` (str): The response generated by the model
    - `reference` (str): The "correct" answer to the question
    - `user_input` (str): The user input question
    If the above data is not in the appropriate column, you can specify different column
    names for these fields using the parameters `retrieved_contexts_column` and `response_column`.

    For example, if your dataset has this data stored in different columns, you can
    pass the following parameters:
    ```python
    {
        "retrieved_contexts_column": "context_info",
        "response_column": "my_answer_col",
        "reference_column": "reference",
        "user_input_column": "user_input",
    }
    ```

    If the data is stored as a dictionary in another column, specify the column and key
    like this:
    ```python
    pred_col = dataset.prediction_column(model)
    params = {
        "reference_column": "reference",
        "retrieved_contexts_column": f"{pred_col}.retrieved_contexts",
        "response_column": f"{pred_col}.response",
        "user_input_column": "user_input",
    }
    ```

    For more complex situations, you can use a function to extract the data:
    ```python
    pred_col = dataset.prediction_column(model)
    params = {
        "reference_column": "reference",
        "retrieved_contexts_column": lambda row: [row[pred_col]["context_message"]],
        "response_column": lambda row: "\\n\\n".join(row[pred_col]["messages"]),
        "user_input_column": "user_input",
    }

    ### Signs of High Risk

    - High noise sensitivity scores across multiple samples.
    - Significant deviation between mean and median noise sensitivity scores.
    - High standard deviation indicating inconsistency in the model's performance.

    ### Strengths

    - Provides a quantitative measure of how well the LLM handles noisy or irrelevant context.
    - Easy integration and configuration using column parameters.
    - Utilizes both histogram and box plot visualizations to analyze score distribution.

    ### Limitations

    - Requires accurate ground truth that aligns with the generated answers.
    - Assumes the context provided is sufficiently granular to assess noise sensitivity.
    - Primarily applicable to tasks like text QA, text generation, and text summarization where contextual relevance is
    critical.
    """
    warnings.filterwarnings(
        "ignore",
        category=FutureWarning,
        message="promote has been superseded by promote_options='default'.",
    )

    if focus not in VALID_FOCUS_VALUES:
        raise ValueError(
            f"Invalid focus parameter: '{focus}'. "
            f"Must be one of: {VALID_FOCUS_VALUES}"
        )

    required_columns = {
        "response": response_column,
        "retrieved_contexts": retrieved_contexts_column,
        "reference": reference_column,
        "user_input": user_input_column,
    }

    df = get_renamed_columns(dataset._df, required_columns)

    result_df = evaluate(
        Dataset.from_pandas(df),
        metrics=[noise_sensitivity(focus=focus)],
        **get_ragas_config(judge_llm, judge_embeddings),
    ).to_pandas()

    score_column = f"noise_sensitivity_{focus}"

    fig_histogram = px.histogram(
        x=result_df[score_column].to_list(),
        nbins=10,
        title=f"Noise Sensitivity ({focus})",
    )
    fig_box = px.box(
        x=result_df[score_column].to_list(),
        title=f"Noise Sensitivity Distribution ({focus})",
    )

    return (
        {
            "Aggregate Scores": [
                {
                    "Mean Score": result_df[score_column].mean(),
                    "Median Score": result_df[score_column].median(),
                    "Max Score": result_df[score_column].max(),
                    "Min Score": result_df[score_column].min(),
                    "Standard Deviation": result_df[score_column].std(),
                    "Count": result_df.shape[0],
                }
            ],
        },
        fig_histogram,
        fig_box,
        RawData(evaluation_results=result_df, dataset=dataset.input_id),
    )
