from __future__ import annotations

from pathlib import Path

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


def choose_f1_threshold(labels: np.ndarray, probabilities: np.ndarray) -> float:
    """Choose the validation threshold with the best F1 score."""
    candidates = np.linspace(0.20, 0.80, 121)
    scores = np.asarray(
        [f1_score(labels, probabilities >= value, zero_division=0) for value in candidates]
    )
    # Several thresholds often tie on a small validation set. Choosing the one
    # closest to the conventional 0.5 avoids an arbitrary low-threshold bias.
    tied = np.flatnonzero(np.isclose(scores, scores.max()))
    best_index = tied[np.argmin(np.abs(candidates[tied] - 0.5))]
    return float(candidates[best_index])


def classification_metrics(
    labels: np.ndarray, probabilities: np.ndarray, threshold: float
) -> dict[str, float | int]:
    predictions = (probabilities >= threshold).astype(int)
    tn, fp, fn, tp = confusion_matrix(labels, predictions, labels=[0, 1]).ravel()
    return {
        "accuracy": float(accuracy_score(labels, predictions)),
        "precision": float(precision_score(labels, predictions, zero_division=0)),
        "recall": float(recall_score(labels, predictions, zero_division=0)),
        "specificity": float(tn / max(tn + fp, 1)),
        "f1": float(f1_score(labels, predictions, zero_division=0)),
        "roc_auc": float(roc_auc_score(labels, probabilities))
        if len(np.unique(labels)) == 2
        else float("nan"),
        "threshold": float(threshold),
        "tn": int(tn),
        "fp": int(fp),
        "fn": int(fn),
        "tp": int(tp),
        "n_samples": int(len(labels)),
    }


def save_evaluation_plots(
    labels: np.ndarray,
    probabilities: np.ndarray,
    threshold: float,
    output_dir: str | Path,
    title: str,
) -> None:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    predictions = (probabilities >= threshold).astype(int)
    matrix = confusion_matrix(labels, predictions, labels=[0, 1])

    plt.figure(figsize=(5, 4))
    sns.heatmap(
        matrix,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=["Binh thuong", "Te nga"],
        yticklabels=["Binh thuong", "Te nga"],
    )
    plt.xlabel("Du doan")
    plt.ylabel("Thuc te")
    plt.title(f"Ma tran nham lan - {title}")
    plt.tight_layout()
    plt.savefig(output_dir / "confusion_matrix.png", dpi=180)
    plt.close()

    if len(np.unique(labels)) == 2:
        fpr, tpr, _ = roc_curve(labels, probabilities)
        auc = roc_auc_score(labels, probabilities)
        plt.figure(figsize=(5, 4))
        plt.plot(fpr, tpr, label=f"AUC = {auc:.3f}")
        plt.plot([0, 1], [0, 1], "--", color="gray")
        plt.xlabel("False Positive Rate")
        plt.ylabel("True Positive Rate")
        plt.title(f"ROC - {title}")
        plt.legend()
        plt.tight_layout()
        plt.savefig(output_dir / "roc_curve.png", dpi=180)
        plt.close()


def save_predictions(
    path: str | Path,
    labels: np.ndarray,
    probabilities: np.ndarray,
    sources: np.ndarray,
    threshold: float,
) -> None:
    pd.DataFrame(
        {
            "source": sources,
            "label": labels.astype(int),
            "probability_fall": probabilities,
            "prediction": (probabilities >= threshold).astype(int),
        }
    ).to_csv(path, index=False)