# 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.vm_models import VMDataset, VMModel


@tags("nlp", "text_data", "visualization")
@tasks("text_classification", "text_summarization")
def TokenDisparity(
    dataset: VMDataset, model: VMModel
) -> Tuple[pd.DataFrame, go.Figure, RawData]:
    """
    Evaluates the token disparity between reference and generated texts, visualizing the results through histograms and
    bar charts, alongside compiling a comprehensive table of descriptive statistics for token counts.

    ### Purpose

    The Token Disparity test aims to assess the difference in the number of tokens between reference texts and texts
    generated by the model. Understanding token disparity is essential for evaluating how well the generated content
    matches the expected length and richness of the reference texts.

    ### Test Mechanism

    The test extracts true and predicted values from the dataset and model. It computes the number of tokens in each
    reference and generated text. The results are visualized using histograms and bar charts to display the
    distribution of token counts. Additionally, a table of descriptive statistics, including the mean, median, standard
    deviation, minimum, and maximum token counts, is compiled to provide a detailed summary of token usage.

    ### Signs of High Risk

    - Significant disparity in token counts between reference and generated texts could indicate issues with text
    generation quality, such as verbosity or lack of detail.
    - Consistently low token counts in generated texts compared to references might suggest that the model is producing
    incomplete or overly concise outputs.

    ### Strengths

    - Provides a simple yet effective evaluation of text length and token usage.
    - Visual representations (histograms and bar charts) make it easier to interpret the distribution and trends of
    token counts.
    - Descriptive statistics offer a concise summary of the model's performance in generating texts of appropriate
    length.

    ### Limitations

    - Token counts alone do not provide a complete assessment of text quality 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)

    # Calculate token counts
    token_counts_true = [len(text.split()) for text in y_true]
    token_counts_pred = [len(text.split()) for text in y_pred]

    # Create a dataframe for reference and generated token counts
    df = pd.DataFrame(
        {"reference_tokens": token_counts_true, "generated_tokens": token_counts_pred}
    )

    figures = []

    # Generate histograms and bar charts for reference and generated token counts
    token_types = ["reference_tokens", "generated_tokens"]
    token_names = ["Reference Tokens", "Generated Tokens"]

    for token_type, token_name in zip(token_types, token_names):
        # Histogram
        hist_fig = go.Figure(data=[go.Histogram(x=df[token_type])])
        hist_fig.update_layout(
            title=f"{token_name} Histogram",
            xaxis_title=token_name,
            yaxis_title="Count",
        )
        figures.append(hist_fig)

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

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

    # Rename columns for clarity
    stats_df.index = stats_df.index.map(
        {"reference_tokens": "Reference Tokens", "generated_tokens": "Generated Tokens"}
    )

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

    return (
        result_df,
        *figures,
        RawData(token_counts_df=df, model=model.input_id, dataset=dataset.input_id),
    )
