"""
Model evaluation module:
Computes comprehensive multi-class metrics, security-oriented performance indices,
confusion matrices, and visualization plots.
"""

import os
import matplotlib
matplotlib.use("Agg")  # Headless backend
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
from typing import Dict, Any, List
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    classification_report, confusion_matrix, roc_auc_score
)

from src.data_loader import CLASS_NAMES, ID_TO_CLASS


def evaluate_predictions(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    y_proba: np.ndarray = None,
    dataset_name: str = "validation"
) -> Dict[str, Any]:
    """
    Computes rigorous evaluation metrics with a focus on security metrics.
    For email threats, false negatives on 'Suspicious' can lead to breaches.
    """
    acc = accuracy_score(y_true, y_pred)
    macro_p = precision_score(y_true, y_pred, average="macro", zero_division=0)
    macro_r = recall_score(y_true, y_pred, average="macro", zero_division=0)
    macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)

    weighted_p = precision_score(y_true, y_pred, average="weighted", zero_division=0)
    weighted_r = recall_score(y_true, y_pred, average="weighted", zero_division=0)
    weighted_f1 = f1_score(y_true, y_pred, average="weighted", zero_division=0)

    # Per-class metrics
    class_p = precision_score(y_true, y_pred, average=None, zero_division=0)
    class_r = recall_score(y_true, y_pred, average=None, zero_division=0)
    class_f1 = f1_score(y_true, y_pred, average=None, zero_division=0)

    cm = confusion_matrix(y_true, y_pred)
    report_str = classification_report(
        y_true, y_pred, target_names=CLASS_NAMES, digits=4, zero_division=0
    )

    metrics = {
        f"{dataset_name}_accuracy": float(acc),
        f"{dataset_name}_macro_precision": float(macro_p),
        f"{dataset_name}_macro_recall": float(macro_r),
        f"{dataset_name}_macro_f1": float(macro_f1),
        f"{dataset_name}_weighted_f1": float(weighted_f1),
    }

    # Add per-class specific metrics
    for idx, name in enumerate(CLASS_NAMES):
        clean_name = name.lower()
        metrics[f"{dataset_name}_{clean_name}_precision"] = float(class_p[idx])
        metrics[f"{dataset_name}_{clean_name}_recall"] = float(class_r[idx])
        metrics[f"{dataset_name}_{clean_name}_f1"] = float(class_f1[idx])

    # Security false negative rate for Suspicious class (class index 2)
    suspicious_idx = 2
    suspicious_fnr = 1.0 - float(class_r[suspicious_idx])
    metrics[f"{dataset_name}_suspicious_false_negative_rate"] = float(suspicious_fnr)

    # Multi-class ROC-AUC if probabilities provided
    if y_proba is not None:
        try:
            auc = roc_auc_score(y_true, y_proba, multi_class="ovr", average="macro")
            metrics[f"{dataset_name}_roc_auc_ovr"] = float(auc)
        except Exception:
            metrics[f"{dataset_name}_roc_auc_ovr"] = 0.0

    return {
        "metrics": metrics,
        "confusion_matrix": cm,
        "classification_report": report_str
    }


def plot_confusion_matrix(
    cm: np.ndarray,
    model_name: str,
    output_path: str
) -> str:
    """
    Renders and saves a high-resolution confusion matrix heatmap.
    """
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    plt.figure(figsize=(7, 6))
    
    # Normalize for percentage annotations
    cm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    annotations = np.empty_like(cm, dtype=object)
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            annotations[i, j] = f"{cm[i, j]}\n({cm_norm[i, j]*100:.1f}%)"

    sns.heatmap(
        cm,
        annot=annotations,
        fmt="",
        cmap="Blues",
        xticklabels=CLASS_NAMES,
        yticklabels=CLASS_NAMES,
        cbar=True,
        annot_kws={"size": 11, "weight": "bold"}
    )
    plt.title(f"Confusion Matrix: {model_name}", fontsize=13, pad=12, weight="bold")
    plt.xlabel("Predicted Class", fontsize=11, weight="bold")
    plt.ylabel("True Class", fontsize=11, weight="bold")
    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close()
    return output_path


def plot_metric_comparisons(
    comparison_df: pd.DataFrame,
    output_path: str
) -> str:
    """
    Plots a multi-bar chart comparing model performances across key metrics.
    """
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    plt.figure(figsize=(10, 6))
    
    # Melt dataframe for seaborn barplot
    metrics_to_plot = ["Macro F1", "Suspicious Recall", "Accuracy", "Weighted F1"]
    melted = comparison_df.melt(
        id_vars=["Model"],
        value_vars=[m for m in metrics_to_plot if m in comparison_df.columns],
        var_name="Metric",
        value_name="Score"
    )

    sns.barplot(
        data=melted,
        x="Metric",
        y="Score",
        hue="Model",
        palette="crest"
    )
    plt.title("Model Performance Benchmark Comparison", fontsize=13, weight="bold", pad=12)
    plt.ylim(0.0, 1.05)
    plt.ylabel("Score (0.0 - 1.0)", fontsize=11, weight="bold")
    plt.xlabel("Evaluation Metric", fontsize=11, weight="bold")
    plt.legend(loc="lower right")
    plt.grid(axis="y", linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close()
    return output_path