"""
Evaluation utilities for Gherbal models.
"""

import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix


def evaluate(model, eval_df):
    """
    Evaluate a Gherbal model on a labeled DataFrame.

    Args:
        model: A Gherbal model instance (must have a .predict() method).
        eval_df: DataFrame with 'lang' and 'text' columns.

    Returns:
        Tuple of (accuracy, report_str, confusion_matrix, present_labels, filtered_df).
    """
    if not isinstance(eval_df, pd.DataFrame):
        eval_df = pd.DataFrame(eval_df)

    y_pred = []
    for text in tqdm(eval_df["text"], desc="Evaluating"):
        y_pred.append(model.predict(text.replace("\n", " ")))
    y_pred_labels = [pred[0][0] for pred in y_pred]

    predicted_labels = set(y_pred_labels)
    true_labels = set(eval_df["lang"])
    known_classes = predicted_labels.intersection(true_labels)

    if not known_classes:
        raise ValueError(
            "No matching classes found between predictions and evaluation data"
        )

    eval_df_filtered = eval_df[eval_df["lang"].isin(known_classes)].copy()
    y_pred_filtered = [
        label
        for i, label in enumerate(y_pred_labels)
        if eval_df.iloc[i]["lang"] in known_classes
    ]

    if len(eval_df_filtered) == 0:
        raise ValueError("No known classes found in evaluation data")

    present_labels = sorted(known_classes)
    acc = accuracy_score(eval_df_filtered["lang"], y_pred_filtered)

    report_dict = classification_report(
        eval_df_filtered["lang"],
        y_pred_filtered,
        labels=present_labels,
        target_names=present_labels,
        zero_division=0,
        output_dict=True,
    )

    # Convert support values to integers
    for label in report_dict:
        if isinstance(report_dict[label], dict) and "support" in report_dict[label]:
            report_dict[label]["support"] = int(report_dict[label]["support"])

    report_str = _format_report(report_dict, present_labels, acc)
    cm = confusion_matrix(
        eval_df_filtered["lang"], y_pred_filtered, labels=present_labels
    )

    return acc, report_str, cm, present_labels, eval_df_filtered


def report(model, eval_df, silent=False):
    """
    Evaluate model and print full report with confused pairs.

    Args:
        model: Gherbal model instance.
        eval_df: DataFrame with 'lang' and 'text' columns.
        silent: If True, suppress printing.

    Returns:
        Tuple of (accuracy, filtered_df).
    """
    acc, report_str, cm, present_labels, eval_df_filtered = evaluate(model, eval_df)

    if not silent:
        print(f"Accuracy: {acc:.4f}")
        print("\nClassification Report:")
        print(report_str)
        print("\nConfusion Matrix:")
        print(cm)
        _print_confused_pairs(cm, present_labels)

    return acc, eval_df_filtered


def get_confused_pairs(cm, labels, top_k=10):
    """
    Extract the most confused label pairs from a confusion matrix.

    Args:
        cm: Confusion matrix (numpy array).
        labels: List of label strings matching cm dimensions.
        top_k: Number of top confused pairs to return.

    Returns:
        List of (label_a, label_b, total_confused, description_str) tuples.
    """
    pairs = []
    for i in range(len(labels)):
        for j in range(i + 1, len(labels)):
            pairs.append((
                labels[i],
                labels[j],
                cm[i, j] + cm[j, i],
                f"{labels[i]}->{labels[j]}: {cm[i, j]}, "
                f"{labels[j]}->{labels[i]}: {cm[j, i]}",
            ))
    pairs.sort(key=lambda x: x[2], reverse=True)
    return pairs[:top_k]


def plot_stats(eval_df, cm, present_labels, title="Evaluation Statistics"):
    """
    Plot class distribution and confusion matrix heatmap.

    Args:
        eval_df: DataFrame with a 'lang' column.
        cm: Confusion matrix.
        present_labels: List of label strings.
        title: Plot title prefix.
    """
    import matplotlib.pyplot as plt
    import seaborn as sns
    from matplotlib.colors import LogNorm

    if not isinstance(eval_df, pd.DataFrame):
        eval_df = pd.DataFrame(eval_df)

    class_counts = eval_df["lang"].value_counts()

    plt.figure(figsize=(15, 8))
    sns.barplot(x=class_counts.index, y=class_counts.values)
    plt.title(f"Class Occurrences in {title}")
    plt.xlabel("Class")
    plt.ylabel("Count")
    plt.xticks(rotation=90)
    plt.tight_layout()
    plt.show()

    plt.figure(figsize=(15, 15))
    sns.heatmap(
        cm,
        annot=False,
        cmap="YlOrRd",
        xticklabels=present_labels,
        yticklabels=present_labels,
        norm=LogNorm(),
    )
    plt.title(f"Confusion Matrix Heatmap ({title}) - Log Scale")
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.tight_layout()
    plt.show()


def _format_report(report_dict, present_labels, accuracy):
    """Format a classification report dict into a readable string."""
    lines = ["               precision    recall  f1-score   support\n"]
    for label in present_labels:
        m = report_dict[label]
        lines.append(
            f"{label:>20} {m['precision']:>10.2f} {m['recall']:>10.2f} "
            f"{m['f1-score']:>10.2f} {m['support']:>10d}"
        )
    total = np.sum([report_dict[l]["support"] for l in present_labels])
    lines.append(
        f"\n{'accuracy':>20} {'':<10} {'':<10} {accuracy:>10.2f} {total:>10d}"
    )
    for avg in ("macro avg", "weighted avg"):
        m = report_dict[avg]
        lines.append(
            f"{avg:>20} {m['precision']:>10.2f} {m['recall']:>10.2f} "
            f"{m['f1-score']:>10.2f} {total:>10d}"
        )
    return "\n".join(lines)


def _print_confused_pairs(cm, labels, top_k=10):
    """Print top confused pairs."""
    pairs = get_confused_pairs(cm, labels, top_k)
    print(f"\nTop {top_k} Most Confused Pairs:")
    for pair in pairs:
        print(pair[3])
