"""Visualization utilities for XAI explanations."""

import matplotlib.pyplot as plt
import numpy as np


def plot_local_explanation(explanation_dict):
    """
    Creates a bar plot for a local explanation.

    Args:
        explanation_dict (dict): The standardized output from an explainer.
    """
    method = explanation_dict.get("method", "Explanation")
    scores = explanation_dict["importance_scores"]
    feature_names = explanation_dict["feature_names"]

    sorted_indices = np.argsort(np.abs(scores))

    plt.figure(figsize=(8, len(feature_names) * 0.4))
    plt.barh(
        np.array(feature_names)[sorted_indices], np.array(scores)[sorted_indices]
    )
    plt.xlabel("Contribution to Prediction")
    plt.title(f"Local Feature Importance ({method})")
    plt.tight_layout()
    plt.show()


def save_local_explanation(explanation_dict, output_path, dpi=100, top_n=None):
    """
    Saves a bar plot for a local explanation to a file.

    Args:
        explanation_dict (dict): The standardized output from an explainer.
        output_path (str): Path where the image will be saved.
        dpi (int): Resolution in dots per inch. Defaults to 100.
        top_n (int, optional): If provided, only show the top N features by
            absolute importance. Defaults to None (show all).

    Returns:
        str: The path where the image was saved.
    """
    method = explanation_dict.get("method", "Explanation")
    scores = np.array(explanation_dict["importance_scores"])
    feature_names = np.array(explanation_dict["feature_names"])

    # Sort by absolute importance
    sorted_indices = np.argsort(np.abs(scores))

    # Apply top_n filter if specified
    if top_n is not None and top_n < len(sorted_indices):
        sorted_indices = sorted_indices[-top_n:]

    sorted_features = feature_names[sorted_indices]
    sorted_scores = scores[sorted_indices]

    fig, ax = plt.subplots(figsize=(8, len(sorted_features) * 0.4))
    ax.barh(sorted_features, sorted_scores)
    ax.set_xlabel("Contribution to Prediction")
    ax.set_title(f"Local Feature Importance ({method})")
    plt.tight_layout()
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)

    return output_path
