from __future__ import annotations

from pathlib import Path
import numpy as np


def confusion_matrix_data(y_true, y_pred) -> np.ndarray:
    """Return 2x2 confusion matrix for binary labels 0/1."""
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    tn = int(((y_true == 0) & (y_pred == 0)).sum())
    fp = int(((y_true == 0) & (y_pred == 1)).sum())
    fn = int(((y_true == 1) & (y_pred == 0)).sum())
    tp = int(((y_true == 1) & (y_pred == 1)).sum())
    return np.array([[tn, fp], [fn, tp]])


def roc_curve_data(y_true, y_score, thresholds=None):
    """Return (threshold, fpr, tpr) points for binary ROC curve."""
    y_true = np.asarray(y_true)
    y_score = np.asarray(y_score)
    thresholds = (
        np.asarray(thresholds) if thresholds is not None else np.linspace(0.0, 1.0, 101)
    )
    rows = []
    for t in thresholds:
        y_pred = (y_score >= t).astype(int)
        cm = confusion_matrix_data(y_true, y_pred)
        tn, fp, fn, tp = cm.ravel()
        fpr = fp / (fp + tn) if (fp + tn) else 0.0
        tpr = tp / (tp + fn) if (tp + fn) else 0.0
        rows.append((float(t), float(fpr), float(tpr)))
    return rows


def pr_curve_data(y_true, y_score, thresholds=None):
    """Return (threshold, precision, recall) points for binary PR curve."""
    y_true = np.asarray(y_true)
    y_score = np.asarray(y_score)
    thresholds = (
        np.asarray(thresholds) if thresholds is not None else np.linspace(0.0, 1.0, 101)
    )
    rows = []
    for t in thresholds:
        y_pred = (y_score >= t).astype(int)
        cm = confusion_matrix_data(y_true, y_pred)
        _, fp, fn, tp = cm.ravel()
        precision = tp / (tp + fp) if (tp + fp) else 0.0
        recall = tp / (tp + fn) if (tp + fn) else 0.0
        rows.append((float(t), float(precision), float(recall)))
    return rows


def maybe_plot_curve(points, x_label: str, y_label: str, title: str):
    """Plot curve if matplotlib is installed; otherwise return False."""
    try:
        import matplotlib.pyplot as plt  # type: ignore
    except Exception:
        return False
    x = [p[1] for p in points]
    y = [p[2] for p in points]
    plt.figure()
    plt.plot(x, y)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.title(title)
    plt.close()
    return True


def save_curve_plot(
    points,
    x_label: str,
    y_label: str,
    title: str,
    out_path: str | Path,
) -> bool:
    """Save curve plot to path if matplotlib is available."""
    try:
        import matplotlib.pyplot as plt  # type: ignore
    except Exception:
        return False
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    x = [p[1] for p in points]
    y = [p[2] for p in points]
    plt.figure()
    plt.plot(x, y)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.title(title)
    plt.tight_layout()
    plt.savefig(out)
    plt.close()
    return True
