utils.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    multilabel_confusion_matrix,
    roc_auc_score,
)

from config import CLASS_DESCRIPTIONS, LEGAL_WARNING, THRESHOLD


def save_json(path: str | Path, payload: dict[str, Any]) -> None:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")


def load_json(path: str | Path, default: Any = None) -> Any:
    source = Path(path)
    if not source.exists():
        return default
    return json.loads(source.read_text(encoding="utf-8"))


def flatten_history(phases: dict[str, dict[str, list[float]]]) -> dict[str, list[float]]:
    flat: dict[str, list[float]] = {}
    for phase_name in ("phase1", "phase2"):
        phase = phases.get(phase_name, {})
        for metric_name, values in phase.items():
            flat.setdefault(metric_name, [])
            flat[metric_name].extend(values)
    return flat


def plot_training_history(history_payload: dict[str, Any], output_path: str | Path) -> None:
    flat = flatten_history(history_payload)
    if not flat:
        return

    metrics_to_plot = [
        ("loss", "val_loss", "Pérdida"),
        ("acc", "val_acc", "Exactitud"),
        ("auc", "val_auc", "AUC"),
    ]

    plt.figure(figsize=(16, 4))
    for index, (train_metric, val_metric, title) in enumerate(metrics_to_plot, start=1):
        plt.subplot(1, 3, index)
        if train_metric in flat:
            plt.plot(flat[train_metric], label=train_metric)
        if val_metric in flat:
            plt.plot(flat[val_metric], label=val_metric)
        plt.title(title)
        plt.xlabel("Época")
        plt.grid(alpha=0.3)
        plt.legend()
    plt.tight_layout()
    plt.savefig(output_path, dpi=180, bbox_inches="tight")
    plt.close()


def _safe_macro_auc(y_true: np.ndarray, y_prob: np.ndarray, multilabel: bool) -> float | None:
    try:
        if multilabel:
            valid_scores = []
            for index in range(y_true.shape[1]):
                column = y_true[:, index]
                if len(np.unique(column)) < 2:
                    continue
                valid_scores.append(roc_auc_score(column, y_prob[:, index]))
            if not valid_scores:
                return None
            return float(np.mean(valid_scores))
        return float(roc_auc_score(y_true, y_prob, multi_class="ovr", average="macro"))
    except Exception:
        return None


def evaluate_predictions(
    y_true: np.ndarray,
    y_prob: np.ndarray,
    class_names: list[str],
    task_type: str,
    threshold: float = THRESHOLD,
) -> dict[str, Any]:
    multilabel = task_type == "multilabel"

    if multilabel:
        y_pred = (y_prob >= threshold).astype(int)
        exact_match = float(accuracy_score(y_true, y_pred))
        macro_f1 = float(f1_score(y_true, y_pred, average="macro", zero_division=0))
        micro_f1 = float(f1_score(y_true, y_pred, average="micro", zero_division=0))
        auc_value = _safe_macro_auc(y_true, y_prob, multilabel=True)
        per_class = {}
        for index, class_name in enumerate(class_names):
            report = classification_report(
                y_true[:, index],
                y_pred[:, index],
                output_dict=True,
                zero_division=0,
            )
            per_class[class_name] = report
        confusion_payload = multilabel_confusion_matrix(y_true, y_pred).tolist()
        return {
            "task_type": task_type,
            "accuracy": exact_match,
            "macro_f1": macro_f1,
            "micro_f1": micro_f1,
            "auc_roc": auc_value,
            "classification_report": per_class,
            "confusion_matrix": confusion_payload,
        }

    y_true_idx = np.argmax(y_true, axis=1)
    y_pred_idx = np.argmax(y_prob, axis=1)
    accuracy = float(accuracy_score(y_true_idx, y_pred_idx))
    macro_f1 = float(f1_score(y_true_idx, y_pred_idx, average="macro", zero_division=0))
    auc_value = _safe_macro_auc(y_true, y_prob, multilabel=False)
    report = classification_report(
        y_true_idx,
        y_pred_idx,
        target_names=class_names,
        output_dict=True,
        zero_division=0,
    )
    cm = confusion_matrix(y_true_idx, y_pred_idx).tolist()
    return {
        "task_type": task_type,
        "accuracy": accuracy,
        "macro_f1": macro_f1,
        "micro_f1": macro_f1,
        "auc_roc": auc_value,
        "classification_report": report,
        "confusion_matrix": cm,
    }


def plot_confusion_figure(
    y_true: np.ndarray,
    y_prob: np.ndarray,
    class_names: list[str],
    task_type: str,
    output_path: str | Path,
    threshold: float = THRESHOLD,
) -> None:
    multilabel = task_type == "multilabel"

    if multilabel:
        y_pred = (y_prob >= threshold).astype(int)
        matrices = multilabel_confusion_matrix(y_true, y_pred)
        columns = 3
        rows = int(np.ceil(len(class_names) / columns))
        plt.figure(figsize=(5 * columns, 4 * rows))
        for index, class_name in enumerate(class_names, start=1):
            plt.subplot(rows, columns, index)
            sns.heatmap(
                matrices[index - 1],
                annot=True,
                fmt="d",
                cmap="Blues",
                cbar=False,
                xticklabels=["Negativo", "Positivo"],
                yticklabels=["Negativo", "Positivo"],
            )
            plt.title(class_name)
            plt.xlabel("Predicción")
            plt.ylabel("Real")
        plt.tight_layout()
        plt.savefig(output_path, dpi=180, bbox_inches="tight")
        plt.close()
        return

    y_true_idx = np.argmax(y_true, axis=1)
    y_pred_idx = np.argmax(y_prob, axis=1)
    cm = confusion_matrix(y_true_idx, y_pred_idx)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names)
    plt.xlabel("Predicción")
    plt.ylabel("Real")
    plt.title("Matriz de confusión")
    plt.tight_layout()
    plt.savefig(output_path, dpi=180, bbox_inches="tight")
    plt.close()


def probabilities_to_table(class_names: list[str], probabilities: np.ndarray) -> pd.DataFrame:
    rows = [{"Clase": name, "Probabilidad": float(prob)} for name, prob in zip(class_names, probabilities)]
    frame = pd.DataFrame(rows)
    return frame.sort_values("Probabilidad", ascending=False).reset_index(drop=True)


def format_diagnosis_message(
    class_names: list[str],
    probabilities: np.ndarray,
    task_type: str,
    threshold: float = THRESHOLD,
) -> str:
    if task_type == "multilabel":
        positives = [(name, float(prob)) for name, prob in zip(class_names, probabilities) if prob >= threshold]
        positives = sorted(positives, key=lambda item: item[1], reverse=True)

        if positives and positives[0][0] == "normal":
            return "✅ Radiografía NORMAL — Sin hallazgos patológicos en las clases objetivo del modelo."

        if positives:
            top_name, top_prob = positives[0]
            description = CLASS_DESCRIPTIONS.get(top_name, "Hallazgo compatible con la clase predicha.")
            return (
                f"⚠️ DIAGNÓSTICO PROBABLE: {top_name}\n"
                f"Confianza: {top_prob * 100:.2f}%\n"
                f"¿Qué significa? {description}"
            )

        top_index = int(np.argmax(probabilities))
        top_name = class_names[top_index]
        top_prob = float(probabilities[top_index])
        description = CLASS_DESCRIPTIONS.get(top_name, "Hallazgo compatible con la clase predicha.")
        return (
            f"⚠️ Hallazgo de baja confianza. Clase más probable: {top_name}\n"
            f"Confianza: {top_prob * 100:.2f}%\n"
            f"¿Qué significa? {description}"
        )

    top_index = int(np.argmax(probabilities))
    top_name = class_names[top_index]
    top_prob = float(probabilities[top_index])
    if top_name == "normal":
        return "✅ Radiografía NORMAL — Sin hallazgos patológicos."
    description = CLASS_DESCRIPTIONS.get(top_name, "Hallazgo compatible con la clase predicha.")
    return (
        f"⚠️ DIAGNÓSTICO PROBABLE: {top_name}\n"
        f"Confianza: {top_prob * 100:.2f}%\n"
        f"¿Qué significa? {description}"
    )


def legal_warning_text() -> str:
    return LEGAL_WARNING