# 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.tests.utils import validate_prediction
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 BleuScore. "
            "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 BleuScore(
    dataset: VMDataset, model: VMModel
) -> Tuple[pd.DataFrame, go.Figure, RawData]:
    """
    Evaluates the quality of machine-generated text using BLEU metrics and visualizes the results through histograms
    and bar charts, alongside compiling a comprehensive table of descriptive statistics for BLEU scores.

    ### Purpose

    This function is designed to assess the quality of text generated by machine learning models using the BLEU metric.
    BLEU, which stands for Bilingual Evaluation Understudy, is a metric used to evaluate the overlap of n-grams between
    the machine-generated text and reference texts. This evaluation is crucial for tasks such as text summarization,
    machine translation, and text generation, where the goal is to produce text that accurately reflects the content
    and meaning of human-crafted references.

    ### Test Mechanism

    The function starts by extracting the true and predicted values from the provided dataset and model. It then
    initializes the BLEU evaluator. For each pair of true and predicted texts, the function calculates the BLEU scores
    and compiles them into a dataframe. Histograms and bar charts are generated for the BLEU scores to visualize their
    distribution. Additionally, a table of descriptive statistics (mean, median, standard deviation, minimum, and
    maximum) is compiled for the BLEU scores, providing a comprehensive summary of the model's performance.

    ### Signs of High Risk

    - Consistently low BLEU scores could indicate poor quality in the generated text, suggesting that the model fails
    to capture the essential content of the reference texts.
    - Low precision scores might suggest that the generated text contains a lot of redundant or irrelevant information.
    - Low recall scores may indicate that important information from the reference text is being omitted.
    - An imbalanced performance between precision and recall, reflected by a low BLEU score, could signal issues in the
    model's ability to balance informativeness and conciseness.

    ### Strengths

    - Provides a straightforward and widely-used evaluation of text quality through BLEU scores.
    - Visual representations (histograms and bar charts) make it easier to interpret the distribution and trends of the
    scores.
    - Descriptive statistics offer a concise summary of the model's strengths and weaknesses in generating text.

    ### Limitations

    - BLEU metrics primarily focus on n-gram overlap and may not fully capture semantic coherence, fluency, or
    grammatical quality of the text.
    - The evaluation relies on the availability of high-quality reference texts, which may not always be obtainable.
    - While useful for comparison, BLEU scores alone do not provide a complete assessment of a model's performance and
    should be supplemented with other metrics and qualitative analysis.
    """

    # Extract true and predicted values
    y_true = dataset.y
    y_pred = dataset.y_pred(model)

    # Ensure equal lengths and get truncated data if necessary
    y_true, y_pred = validate_prediction(y_true, y_pred)

    bleu = evaluate.load("bleu")

    # Calculate BLEU scores
    score_list = []
    for y_t, y_p in zip(y_true, y_pred):
        # Compute the BLEU score
        score = bleu.compute(predictions=[y_p], references=[[y_t]])
        score_list.append(score["bleu"])

    # Convert scores to a dataframe
    metrics_df = pd.DataFrame(score_list, columns=["BLEU Score"])

    figures = []

    # Histogram for BLEU Score
    hist_fig = go.Figure(data=[go.Histogram(x=metrics_df["BLEU Score"])])
    hist_fig.update_layout(
        title="BLEU Score Histogram",
        xaxis_title="BLEU Score",
        yaxis_title="Count",
    )
    figures.append(hist_fig)

    # Bar Chart for BLEU Score
    bar_fig = go.Figure(data=[go.Bar(x=metrics_df.index, y=metrics_df["BLEU Score"])])
    bar_fig.update_layout(
        title="BLEU Score Bar Chart",
        xaxis_title="Row Index",
        yaxis_title="BLEU Score",
    )
    figures.append(bar_fig)

    # Calculate statistics for BLEU Score
    stats_df = metrics_df.describe().loc[["mean", "50%", "max", "min", "std"]]
    stats_df = stats_df.rename(
        index={
            "mean": "Mean Score",
            "50%": "Median Score",
            "max": "Max Score",
            "min": "Min Score",
            "std": "Standard Deviation",
        }
    ).T
    stats_df["Count"] = len(metrics_df)

    # Create a DataFrame from all collected statistics
    result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})

    return (
        result_df,
        *figures,
        RawData(
            bleu_scores_df=metrics_df, model=model.input_id, dataset=dataset.input_id
        ),
    )
