# 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 BertScore. "
            "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 BertScore(
    dataset: VMDataset,
    model: VMModel,
    evaluation_model="distilbert-base-uncased",
) -> Tuple[pd.DataFrame, go.Figure, RawData]:
    """
    Assesses the quality of machine-generated text using BERTScore metrics and visualizes results through histograms
    and bar charts, alongside compiling a comprehensive table of descriptive statistics.

    ### Purpose

    This function is designed to assess the quality of text generated by machine learning models using BERTScore
    metrics. BERTScore evaluates text generation models' performance by calculating precision, recall, and F1 score
    based on BERT contextual embeddings.

    ### Test Mechanism

    The function starts by extracting the true and predicted values from the provided dataset and model. It then
    initializes the BERTScore evaluator. For each pair of true and predicted texts, the function calculates the
    BERTScore metrics and compiles them into a dataframe. Histograms and bar charts are generated for each BERTScore
    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. The test uses the `evaluation_model` param to specify the
    huggingface model to use for evaluation. `microsoft/deberta-xlarge-mnli` is the best-performing model but is
    very large and may be slow without a GPU. `microsoft/deberta-large-mnli` is a smaller model that is faster to
    run and `distilbert-base-uncased` is much lighter and can run on a CPU but is less accurate.

    ### Signs of High Risk

    - Consistently low scores across BERTScore 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 BERTScore 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

    - BERTScore relies on the contextual embeddings from BERT models, which may not fully capture all nuances of text
    similarity.
    - The evaluation relies on the availability of high-quality reference texts, which may not always be obtainable.
    - While useful for comparison, BERTScore metrics 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)

    bert = evaluate.load("bertscore")

    # Compute the BERT score
    bert_s = bert.compute(
        predictions=y_pred,
        references=y_true,
        lang="en",
        model_type=evaluation_model,
    )

    # Convert scores to a dataframe
    metrics_df = pd.DataFrame(bert_s)
    figures = []

    # Generate histograms and bar charts for each score type
    score_types = ["precision", "recall", "f1"]
    score_names = ["Precision", "Recall", "F1 Score"]

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

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

    # Calculate statistics for each score type
    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)

    # Rename metrics for clarity
    stats_df.index = stats_df.index.map(
        {"precision": "Precision", "recall": "Recall", "f1": "F1 Score"}
    )

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

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