src/visualization/plots.py

"""
Visualization utilities for fraud detection analysis and reporting.

Production-ready plotting functions for:
- Model performance (ROC, PR curves)
- Feature importance
- Fraud distribution analysis
- Monitoring dashboards
"""

import logging
from pathlib import Path
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from sklearn.metrics import auc, precision_recall_curve, roc_curve

logger = logging.getLogger(__name__)

# Consistent styling
plt.style.use("seaborn-v0_8-whitegrid")
COLORS = {"fraud": "#e74c3c", "legit": "#2ecc71", "primary": "#3498db", "secondary": "#9b59b6"}


def plot_roc_curves(
    results: dict[str, tuple[np.ndarray, np.ndarray]],
    save_path: Optional[str] = None,
) -> plt.Figure:
    """Plot ROC curves for multiple models on a single figure.

    Args:
        results: Dict mapping model name to (y_true, y_proba) tuples.
        save_path: Optional path to save the figure.

    Returns:
        Matplotlib figure.
    """
    fig, ax = plt.subplots(figsize=(10, 8))

    for name, (y_true, y_proba) in results.items():
        fpr, tpr, _ = roc_curve(y_true, y_proba)
        roc_auc = auc(fpr, tpr)
        ax.plot(fpr, tpr, linewidth=2, label=f"{name} (AUC = {roc_auc:.4f})")

    ax.plot([0, 1], [0, 1], "k--", linewidth=1, alpha=0.5, label="Random")
    ax.set_xlabel("False Positive Rate", fontsize=12)
    ax.set_ylabel("True Positive Rate", fontsize=12)
    ax.set_title("ROC Curves — Model Comparison", fontsize=14, fontweight="bold")
    ax.legend(loc="lower right", fontsize=11)
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1.02])

    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
        logger.info("ROC curves saved to %s", save_path)
    return fig


def plot_precision_recall_curves(
    results: dict[str, tuple[np.ndarray, np.ndarray]],
    save_path: Optional[str] = None,
) -> plt.Figure:
    """Plot Precision-Recall curves (critical for imbalanced fraud detection).

    Args:
        results: Dict mapping model name to (y_true, y_proba) tuples.
        save_path: Optional path to save the figure.

    Returns:
        Matplotlib figure.
    """
    fig, ax = plt.subplots(figsize=(10, 8))

    for name, (y_true, y_proba) in results.items():
        precision, recall, _ = precision_recall_curve(y_true, y_proba)
        pr_auc = auc(recall, precision)
        ax.plot(recall, precision, linewidth=2, label=f"{name} (AP = {pr_auc:.4f})")

    # Baseline: fraud rate
    baseline = list(results.values())[0][0].mean()
    ax.axhline(y=baseline, color="k", linestyle="--", alpha=0.5, label=f"Baseline ({baseline:.4f})")

    ax.set_xlabel("Recall", fontsize=12)
    ax.set_ylabel("Precision", fontsize=12)
    ax.set_title("Precision-Recall Curves — Model Comparison", fontsize=14, fontweight="bold")
    ax.legend(loc="upper right", fontsize=11)
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1.02])

    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig


def plot_feature_importance(
    importance: dict[str, float],
    top_k: int = 20,
    title: str = "Feature Importance",
    save_path: Optional[str] = None,
) -> plt.Figure:
    """Plot horizontal bar chart of top feature importances.

    Args:
        importance: Dict mapping feature names to importance scores.
        top_k: Number of top features to show.
        title: Chart title.
        save_path: Optional path to save.

    Returns:
        Matplotlib figure.
    """
    sorted_features = sorted(importance.items(), key=lambda x: x[1], reverse=True)[:top_k]
    names, values = zip(*reversed(sorted_features))

    fig, ax = plt.subplots(figsize=(10, max(6, top_k * 0.4)))
    bars = ax.barh(names, values, color=COLORS["primary"], edgecolor="white", linewidth=0.5)
    ax.set_xlabel("Importance Score", fontsize=12)
    ax.set_title(title, fontsize=14, fontweight="bold")

    # Add value labels
    for bar, val in zip(bars, values):
        ax.text(bar.get_width() + max(values) * 0.01, bar.get_y() + bar.get_height() / 2,
                f"{val:.4f}", va="center", fontsize=9)

    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig


def plot_fraud_distribution(
    df,
    amount_col: str = "transaction_amount",
    fraud_col: str = "is_fraud",
    save_path: Optional[str] = None,
) -> plt.Figure:
    """Plot transaction amount distributions for fraud vs legitimate.

    Args:
        df: Transaction DataFrame.
        amount_col: Amount column name.
        fraud_col: Fraud label column name.
        save_path: Optional path to save.

    Returns:
        Matplotlib figure.
    """
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # Amount distribution by class
    ax = axes[0, 0]
    for label, color, name in [(0, COLORS["legit"], "Legitimate"), (1, COLORS["fraud"], "Fraud")]:
        subset = df[df[fraud_col] == label][amount_col]
        ax.hist(subset, bins=50, alpha=0.6, color=color, label=name, density=True)
    ax.set_xlabel("Transaction Amount ($)")
    ax.set_ylabel("Density")
    ax.set_title("Amount Distribution by Class")
    ax.legend()

    # Log-scale amount distribution
    ax = axes[0, 1]
    for label, color, name in [(0, COLORS["legit"], "Legitimate"), (1, COLORS["fraud"], "Fraud")]:
        subset = np.log1p(df[df[fraud_col] == label][amount_col])
        ax.hist(subset, bins=50, alpha=0.6, color=color, label=name, density=True)
    ax.set_xlabel("Log Transaction Amount")
    ax.set_ylabel("Density")
    ax.set_title("Log Amount Distribution by Class")
    ax.legend()

    # Fraud by hour of day
    ax = axes[1, 0]
    if "hour_of_day" in df.columns:
        fraud_by_hour = df.groupby("hour_of_day")[fraud_col].mean() * 100
        ax.bar(fraud_by_hour.index, fraud_by_hour.values, color=COLORS["primary"], alpha=0.8)
        ax.set_xlabel("Hour of Day")
        ax.set_ylabel("Fraud Rate (%)")
        ax.set_title("Fraud Rate by Hour")
        ax.set_xticks(range(0, 24, 2))

    # Fraud by merchant category
    ax = axes[1, 1]
    if "merchant_category_code" in df.columns:
        fraud_by_mcc = df.groupby("merchant_category_code")[fraud_col].mean().sort_values(ascending=True) * 100
        ax.barh(fraud_by_mcc.index, fraud_by_mcc.values, color=COLORS["secondary"], alpha=0.8)
        ax.set_xlabel("Fraud Rate (%)")
        ax.set_title("Fraud Rate by Merchant Category")

    plt.suptitle("Fraud Distribution Analysis", fontsize=16, fontweight="bold", y=1.02)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig


def plot_monitoring_dashboard(
    drift_results: dict,
    performance_history: list[dict],
    save_path: Optional[str] = None,
) -> plt.Figure:
    """Plot monitoring dashboard with drift and performance trends.

    Args:
        drift_results: PSI results from DistributionDriftDetector.
        performance_history: List of performance evaluation results.
        save_path: Optional path to save.

    Returns:
        Matplotlib figure.
    """
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    # PSI drift scores
    ax = axes[0]
    if drift_results:
        features = list(drift_results.keys())[:15]
        psi_values = [drift_results[f]["psi"] for f in features]
        colors = [
            "#2ecc71" if v < 0.1 else "#f39c12" if v < 0.2 else "#e74c3c"
            for v in psi_values
        ]
        ax.barh(features, psi_values, color=colors)
        ax.axvline(x=0.1, color="orange", linestyle="--", alpha=0.7, label="Warning (0.1)")
        ax.axvline(x=0.2, color="red", linestyle="--", alpha=0.7, label="Critical (0.2)")
        ax.set_xlabel("PSI Score")
        ax.set_title("Feature Drift (PSI)")
        ax.legend()

    # Performance trend
    ax = axes[1]
    if performance_history:
        periods = range(len(performance_history))
        f1_values = [h.get("f1", 0) for h in performance_history]
        auc_values = [h.get("auc_roc", 0) for h in performance_history]
        ax.plot(periods, f1_values, "o-", color=COLORS["primary"], label="F1 Score")
        ax.plot(periods, auc_values, "s-", color=COLORS["secondary"], label="AUC-ROC")
        ax.set_xlabel("Evaluation Period")
        ax.set_ylabel("Score")
        ax.set_title("Model Performance Trend")
        ax.legend()
        ax.set_ylim([0, 1])

    plt.suptitle("Production Monitoring Dashboard", fontsize=14, fontweight="bold")
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig