# 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:
    from rouge import Rouge
except ImportError as e:
    if "rouge" in str(e):
        raise MissingDependencyError(
            "Missing required package `rouge` for RougeScore. "
            "Please run `pip install validmind[nlp]` to use NLP tests",
            required_dependencies=["rouge"],
            extra="nlp",
        ) from e

    raise e


@tags("nlp", "text_data", "visualization")
@tasks("text_classification", "text_summarization")
def RougeScore(
    dataset: VMDataset,
    model: VMModel,
    metric: str = "rouge-1",
) -> Tuple[pd.DataFrame, go.Figure, RawData]:
    """
    Assesses the quality of machine-generated text using ROUGE metrics and visualizes the results to provide
    comprehensive performance insights.

    ### Purpose

    The ROUGE Score test is designed to evaluate the quality of text generated by machine learning models using various
    ROUGE metrics. ROUGE, which stands for Recall-Oriented Understudy for Gisting Evaluation, measures the overlap of
    n-grams, word sequences, and word pairs between machine-generated text and reference texts. This evaluation is
    crucial for tasks like 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 test extracts the true and predicted values from the provided dataset and model. It initializes the ROUGE
    evaluator with the specified metric (e.g., ROUGE-1). For each pair of true and predicted texts, it calculates the
    ROUGE scores and compiles them into a dataframe. Histograms and bar charts are generated for each ROUGE metric
    (Precision, Recall, and F1 Score) to visualize their distribution. Additionally, a table of descriptive statistics
    (mean, median, standard deviation, minimum, and maximum) is compiled for each metric, providing a comprehensive
    summary of the model's performance.

    ### Signs of High Risk

    - Consistently low scores across ROUGE metrics 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 F1 Score, could signal issues in the
    model's ability to balance informativeness and conciseness.

    ### Strengths

    - Provides a multifaceted evaluation of text quality through different ROUGE metrics, offering a detailed view of
    model performance.
    - 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

    - ROUGE 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, ROUGE 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)

    # Initialize Rouge with the specified metric
    rouge = Rouge(metrics=[metric])

    # Calculate ROUGE scores
    score_list = []
    for y_t, y_p in zip(y_true, y_pred):
        scores = rouge.get_scores(y_p, y_t, avg=True)
        score_list.append(scores)

    # Convert scores to a dataframe
    metrics_df = pd.DataFrame(score_list)
    df_scores = pd.DataFrame(metrics_df[metric].tolist())

    # Generate histograms and bar charts for each score type
    score_types = ["p", "r", "f"]
    score_names = ["Precision", "Recall", "F1 Score"]
    figures = []

    for score_type, score_name in zip(score_types, score_names):
        # Histogram
        hist_fig = go.Figure(data=[go.Histogram(x=df_scores[score_type])])
        hist_fig.update_layout(
            title=f"{score_name} Histogram for {metric.upper()}",
            xaxis_title=score_name,
            yaxis_title="Count",
        )
        figures.append(hist_fig)

        # Bar Chart
        bar_fig = go.Figure(data=[go.Bar(x=df_scores.index, y=df_scores[score_type])])
        bar_fig.update_layout(
            title=f"{score_name} Bar Chart for {metric.upper()}",
            xaxis_title="Row Index",
            yaxis_title=score_name,
        )
        figures.append(bar_fig)

    # Calculate statistics for each score type
    stats_df = df_scores.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(df_scores)

    # Rename metrics for clarity
    stats_df.index = stats_df.index.map(
        {"p": "Precision", "r": "Recall", "f": "F1 Score"}
    )

    return (
        pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"}),
        *figures,
        RawData(
            rouge_scores_df=df_scores, model=model.input_id, dataset=dataset.input_id
        ),
    )
