"""Matplotlib/seaborn plot writers for classification and regression insights."""
from pathlib import Path

import numpy as np


def _save_or_return(fig, output_path):
    """Save figure to output_path (PNG) or return None if no path given."""
    if output_path is None:
        return None
    import matplotlib
    matplotlib.use("Agg")
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, bbox_inches="tight", dpi=150)
    return path.resolve()


def plot_confusion_matrix(y_true, y_pred, classes=None, title=None, output_path=None):
    """Seaborn heatmap of the confusion matrix.

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import seaborn as sns
    from sklearn.metrics import confusion_matrix

    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    if classes is None:
        classes = np.unique(np.concatenate([y_true, y_pred]))

    cm = confusion_matrix(y_true, y_pred, labels=classes)

    fig, ax = plt.subplots(figsize=(max(5, len(classes)), max(4, len(classes) - 1)))
    sns.heatmap(
        cm, annot=True, fmt="d", cmap="Blues",
        xticklabels=classes, yticklabels=classes, ax=ax,
    )
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title(title or "Confusion Matrix")
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_roc_curves(y_true, y_proba, classes=None, title=None, output_path=None):
    """One ROC curve per class plus macro-average.

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from sklearn.metrics import roc_curve, auc
    from sklearn.preprocessing import label_binarize

    y_true = np.asarray(y_true)
    y_proba = np.asarray(y_proba)
    if classes is None:
        classes = np.unique(y_true)

    fig, ax = plt.subplots(figsize=(7, 5))

    if len(classes) == 2:
        fpr, tpr, _ = roc_curve(y_true, y_proba[:, 1])
        roc_auc = auc(fpr, tpr)
        ax.plot(fpr, tpr, label=f"ROC (AUC={roc_auc:.3f})")
    else:
        y_bin = label_binarize(y_true, classes=classes)
        all_fpr = np.unique(np.concatenate([
            roc_curve(y_bin[:, i], y_proba[:, i])[0]
            for i in range(len(classes))
        ]))
        mean_tpr = np.zeros_like(all_fpr)

        for i, cls in enumerate(classes):
            fpr, tpr, _ = roc_curve(y_bin[:, i], y_proba[:, i])
            roc_auc = auc(fpr, tpr)
            ax.plot(fpr, tpr, alpha=0.5, label=f"Class {cls} (AUC={roc_auc:.3f})")
            mean_tpr += np.interp(all_fpr, fpr, tpr)

        mean_tpr /= len(classes)
        macro_auc = auc(all_fpr, mean_tpr)
        ax.plot(all_fpr, mean_tpr, "k--", lw=2, label=f"Macro avg (AUC={macro_auc:.3f})")

    ax.plot([0, 1], [0, 1], "k:", lw=1)
    ax.set_xlabel("False Positive Rate")
    ax.set_ylabel("True Positive Rate")
    ax.set_title(title or "ROC Curves")
    ax.legend(loc="lower right", fontsize=8)
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_feature_importances(importances_dict, title=None, output_path=None):
    """Horizontal bar chart of feature importances.

    Parameters
    ----------
    importances_dict : dict
        ``{feature_name: importance}`` (will be sorted descending).

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    items = sorted(importances_dict.items(), key=lambda x: x[1], reverse=True)
    names = [k for k, _ in items]
    values = [v for _, v in items]

    # Limit to top-30 for readability
    max_show = 30
    if len(names) > max_show:
        names = names[:max_show]
        values = values[:max_show]

    fig, ax = plt.subplots(figsize=(7, max(3, len(names) * 0.35)))
    y_pos = np.arange(len(names))
    ax.barh(y_pos, values[::-1])
    ax.set_yticks(y_pos)
    ax.set_yticklabels(names[::-1], fontsize=8)
    ax.set_xlabel("Importance")
    ax.set_title(title or "Feature Importances")
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_actual_vs_predicted(y_true, y_pred, title=None, output_path=None):
    """Scatter plot of actual vs predicted with identity line.

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)

    fig, ax = plt.subplots(figsize=(5, 5))
    ax.scatter(y_true, y_pred, alpha=0.4, s=15, edgecolors="none")
    lo = min(y_true.min(), y_pred.min())
    hi = max(y_true.max(), y_pred.max())
    ax.plot([lo, hi], [lo, hi], "r--", lw=1.5, label="y=x")
    ax.set_xlabel("Actual")
    ax.set_ylabel("Predicted")
    ax.set_title(title or "Actual vs Predicted")
    ax.legend()
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_residuals(y_true, y_pred, title=None, output_path=None):
    """Residual (predicted - actual) vs predicted scatter.

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    residuals = y_pred - y_true

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.scatter(y_pred, residuals, alpha=0.4, s=15, edgecolors="none")
    ax.axhline(0, color="r", linestyle="--", lw=1.5)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Residual (pred - actual)")
    ax.set_title(title or "Residuals vs Predicted")
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_residual_distribution(y_true, y_pred, title=None, output_path=None):
    """Histogram of residuals with KDE overlay.

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from scipy.stats import gaussian_kde

    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    residuals = y_pred - y_true

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(residuals, bins=30, density=True, alpha=0.6, color="steelblue", edgecolor="white")

    try:
        kde = gaussian_kde(residuals)
        x_range = np.linspace(residuals.min(), residuals.max(), 200)
        ax.plot(x_range, kde(x_range), "r-", lw=2, label="KDE")
        ax.legend()
    except Exception:
        pass

    ax.set_xlabel("Residual")
    ax.set_ylabel("Density")
    ax.set_title(title or "Residual Distribution")
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result


def plot_tree_quality_scores(scores_array, title=None, output_path=None):
    """Histogram of per-tree quality scores with mean/std annotation.

    Parameters
    ----------
    scores_array : array-like of shape (n_estimators,)

    Returns
    -------
    Path | None
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    scores = np.asarray(scores_array, dtype=float)
    mean = scores.mean()
    std = scores.std()

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(scores, bins=min(20, len(scores)), edgecolor="white", color="steelblue", alpha=0.8)
    ax.axvline(mean, color="r", linestyle="--", lw=1.5, label=f"Mean={mean:.4f}")
    ax.axvline(mean - std, color="orange", linestyle=":", lw=1, label=f"±1σ ({std:.4f})")
    ax.axvline(mean + std, color="orange", linestyle=":", lw=1)
    ax.set_xlabel("Quality Score")
    ax.set_ylabel("Count")
    ax.set_title(title or "Per-Tree Quality Scores")
    ax.legend(fontsize=9)
    fig.tight_layout()
    result = _save_or_return(fig, output_path)
    plt.close(fig)
    return result
