# 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

from typing import Tuple

import pandas as pd
import plotly.graph_objects as go

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

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


@tags("nlp", "text_data", "visualization")
@tasks("text_classification", "text_summarization")
def ToxicityScore(
    dataset: VMDataset, model: VMModel
) -> Tuple[pd.DataFrame, go.Figure, RawData]:
    """
    Assesses the toxicity levels of texts generated by NLP models to identify and mitigate harmful or offensive content.

    ### Purpose

    The ToxicityScore metric is designed to evaluate the toxicity levels of texts generated by models. This is crucial
    for identifying and mitigating harmful or offensive content in machine-generated texts.

    ### Test Mechanism

    The function starts by extracting the input, true, and predicted values from the provided dataset and model. The
    toxicity score is computed for each text using a preloaded `toxicity` evaluation tool. The scores are compiled into
    dataframes, and histograms and bar charts are generated to visualize the distribution of toxicity scores.
    Additionally, a table of descriptive statistics (mean, median, standard deviation, minimum, and maximum) is
    compiled for the toxicity scores, providing a comprehensive summary of the model's performance.

    ### Signs of High Risk

    - Drastic spikes in toxicity scores indicate potentially toxic content within the associated text segment.
    - Persistent high toxicity scores across multiple texts may suggest systemic issues in the model's text generation
    process.

    ### Strengths

    - Provides a clear evaluation of toxicity levels in generated texts, helping to ensure content safety and
    appropriateness.
    - Visual representations (histograms and bar charts) make it easier to interpret the distribution and trends of
    toxicity scores.
    - Descriptive statistics offer a concise summary of the model's performance in generating non-toxic texts.

    ### Limitations

    - The accuracy of the toxicity scores is contingent upon the underlying `toxicity` tool.
    - The scores provide a broad overview but do not specify which portions or tokens of the text are responsible for
    high toxicity.
    - Supplementary, in-depth analysis might be needed for granular insights.
    """

    # Extract true, predicted, and input values
    y_true = dataset.y
    y_pred = dataset.y_pred(model)
    input_text = dataset.df[dataset.text_column]

    # Load the toxicity evaluation metric
    toxicity = evaluate.load("toxicity")

    # Function to calculate toxicity scores
    def compute_toxicity_scores(texts):
        scores = []
        for text in texts:
            score = toxicity.compute(predictions=[text])
            scores.append(score["toxicity"])
        return scores

    # Calculate toxicity scores for input, true, and predicted texts
    input_toxicity = compute_toxicity_scores(input_text)
    true_toxicity = compute_toxicity_scores(y_true)
    pred_toxicity = compute_toxicity_scores(y_pred)

    # Convert scores to dataframes
    input_df = pd.DataFrame(input_toxicity, columns=["Input Text Toxicity"])
    true_df = pd.DataFrame(true_toxicity, columns=["True Text Toxicity"])
    pred_df = pd.DataFrame(pred_toxicity, columns=["Predicted Text Toxicity"])

    figures = []

    # Function to create histogram and bar chart for toxicity scores
    def create_figures(df, title):
        # Histogram
        hist_fig = go.Figure(data=[go.Histogram(x=df.iloc[:, 0])])
        hist_fig.update_layout(
            title=f"{title} Histogram",
            xaxis_title=title,
            yaxis_title="Count",
        )
        figures.append(hist_fig)

        # Bar Chart
        bar_fig = go.Figure(data=[go.Bar(x=df.index, y=df.iloc[:, 0])])
        bar_fig.update_layout(
            title=f"{title} Bar Chart",
            xaxis_title="Text Instance Index",
            yaxis_title=title,
        )
        figures.append(bar_fig)

    # Create figures for each toxicity score dataframe
    create_figures(input_df, "Input Text Toxicity")
    create_figures(true_df, "True Text Toxicity")
    create_figures(pred_df, "Predicted Text Toxicity")

    # Calculate statistics for each toxicity score dataframe
    def calculate_stats(df):
        stats = df.describe().loc[["mean", "50%", "max", "min", "std"]].T
        stats.columns = [
            "Mean Score",
            "Median Score",
            "Max Score",
            "Min Score",
            "Standard Deviation",
        ]
        stats["Metric"] = df.columns[0]
        stats["Count"] = len(df)
        return stats

    input_stats = calculate_stats(input_df)
    true_stats = calculate_stats(true_df)
    pred_stats = calculate_stats(pred_df)

    # Combine statistics into a single dataframe
    result_df = (
        pd.concat([input_stats, true_stats, pred_stats])
        .reset_index()
        .rename(columns={"index": "Statistic"})
    )
    result_df = result_df[
        [
            "Metric",
            "Mean Score",
            "Median Score",
            "Max Score",
            "Min Score",
            "Standard Deviation",
            "Count",
        ]
    ]

    return (
        result_df,
        *figures,
        RawData(
            input_toxicity_df=input_df,
            true_toxicity_df=true_df,
            pred_toxicity_df=pred_df,
            model=model.input_id,
            dataset=dataset.input_id,
        ),
    )
