"""Visualization utilities for explanations in the MUPPET benchmark framework.

This module provides functions for visualizing explanations generated by XAI
methods across different data modalities (images, tabular data, time series).
It includes utilities for denormalization, overlay plotting, and multi-modal
explanation visualization.

Functions:
    revert_to_original_image: Denormalize images for visualization
    plot_explanation_image: Visualize image explanations with heatmap overlays
    plot_explanation_tabular: Create bar charts for tabular feature importance
    plot_explanation_multichannel_time_series: Visualize time series explanations with color coding
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch


def revert_to_original_image(
    image: np.ndarray,
    normalize_mean: list[float] = [0.485, 0.456, 0.406],
    normalize_std: list[float] = [0.229, 0.224, 0.225],
) -> np.ndarray:
    """Reverts a normalized image to its original pixel values.

    This function takes a normalized image (typically from a deep learning model's
    preprocessing step) and applies the inverse of the normalization formula
    to restore the original image's pixel values. By default, it uses the
    standard mean and standard deviation values from ImageNet. The pixel values are
    clipped to the [0, 1] range after denormalization.

    Args:
        image (np.ndarray): The normalized input image as a NumPy array with
            shape (C, H, W).
        normalize_mean (list[float], optional): The mean values used for
            normalization. Defaults to ImageNet's mean.
        normalize_std (list[float], optional): The standard deviation values
            used for normalization. Defaults to ImageNet's standard deviation.

    Returns:
        np.ndarray: The denormalized image as a NumPy array with pixel values
            in the [0, 1] range.

    Notes:
        The denormalization formula is `image = (image * normalize_std) + normalize_mean`.
    """
    image = np.array(normalize_std) * image + np.array(normalize_mean)
    image = np.clip(image, 0, 1)

    return image


def plot_explanation_image(
    example: torch.Tensor | np.ndarray,
    explanation: torch.Tensor | np.ndarray,
    target: str | None = None,
    normalize_mean: list[float] | None = [0.485, 0.456, 0.406],
    normalize_std: list[float] | None = [0.229, 0.224, 0.225],
    image_title: str = "",
    figure_title: str = "",
    caption: str = "",
    save_path: Path | None = None,
    cmap: str = "jet",
    **kwargs,
) -> None:
    """Visualizes an image alongside a superimposed heatmap.

    This function creates a two-subplot figure. The first subplot displays the
    original image, and the second subplot shows the same image with a
    saliency map (heatmap) overlaid on top. It supports both PyTorch tensors
    and NumPy arrays as input and can denormalize the image if normalization
    parameters are provided. The resulting figure can either be displayed
    or saved to a file.

    Args:
        example (torch.Tensor | np.ndarray): The original input image, with shape
            (C, H, W) for tensors or (H, W, C) for NumPy arrays.
        explanation (torch.Tensor | np.ndarray): The saliency map or heatmap data,
            with shape (H, W).
        normalize_mean (list[float] | None, optional): Mean values for
            denormalization. If `None`, no denormalization is applied.
            Defaults to ImageNet's mean.
        normalize_std (list[float] | None, optional): Standard deviation values
            for denormalization. If `None`, no denormalization is applied.
            Defaults to ImageNet's standard deviation.
        title (str, optional): Title for the original image subplot.
            Defaults to an empty string.
        heatmap_title (str, optional): Title for the heatmap subplot.
            Defaults to an empty string.
        caption (str, optional): A caption to display below the entire figure.
            Defaults to an empty string.
        save_path (pathlib.Path, optional): The file path to save the figure. If `None`,
            the figure is displayed. Defaults to `None`.
        cmap (str, optional): The Matplotlib colormap to use for the heatmap.
            Defaults to 'jet'.
        **kwargs: Additional keyword arguments that can be passed to the underlying
            matplotlib functions, although none are currently used in the function body.

    Returns:
        None: The function either displays the figure or saves it to a file.
    """
    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(6, 3))

    # Convert PyTorch tensors to NumPy arrays and adjust shape for plotting
    if isinstance(example, torch.Tensor):
        # Shape of PyTorch tensor is (C, H, W), Matplotlib expects (H, W, C)
        example = example.detach().cpu().numpy().transpose((1, 2, 0))
    if isinstance(explanation, torch.Tensor):
        explanation = explanation.detach().cpu().numpy()

    # Denormalize the image if normalization parameters are provided
    if normalize_mean is not None and normalize_std is not None:
        example = revert_to_original_image(
            example, normalize_mean, normalize_std
        )

    axes[0].imshow(example)
    axes[0].title.set_text(image_title)
    axes[0].axis("off")

    axes[1].imshow(example)
    axes[1].title.set_text(figure_title)
    axes[1].axis("off")
    t = axes[1].imshow(explanation, cmap=cmap, alpha=0.5)
    fig.colorbar(t, fraction=0.046, pad=0.04)
    fig.text(0.5, 0.02, caption, ha="center")

    if save_path is not None:
        # save image as numpy array file
        np.save(
            save_path.with_stem(f"{save_path.stem}_image").with_suffix(".npy"),
            example,
        )
        # save explanation (heatmap) as numpy array file
        np.save(
            save_path.with_stem(f"{save_path.stem}_explanation").with_suffix(
                ".npy"
            ),
            explanation,
        )
        if target is not None:
            # save target as label in a .txt file
            with open(
                save_path.with_stem(f"{save_path.stem}_target").with_suffix(
                    ".txt"
                ),
                mode="w",
            ) as f:
                f.write(str(target))
        # save explanation (heatmap) figure
        fig.savefig(save_path, bbox_inches="tight")
        plt.close(fig)  # Close the figure to prevent it from displaying
    else:
        fig.show()


def plot_explanation_tabular(
    explanation: torch.Tensor | np.ndarray,
    feature_names: list[str] | np.ndarray | None = None,
    figure_title: str = "Feature Importance",
    save_path: str | None = None,
    **kwargs,
):
    """Plots a horizontal bar chart of feature importance.

    This function visualizes the importance or contribution of features as a horizontal
    bar chart. It's particularly useful for interpreting models like SHAP or LIME for
    tabular data. The features are sorted by their importance values in ascending order.
    The plot can be either displayed interactively or saved to a file.

    Args:
        explanation (torch.Tensor | np.ndarray): A 1D or 2D array-like object containing
            the importance values for each feature. For a 2D input, it will be flattened.
        feature_names (list[str] | np.ndarray | None, optional): A list or array of strings
            that provide the names for each feature. If `None`, the feature indices
            will be used as labels. Defaults to `None`.
        heatmap_title (str, optional): The title for the plot. Defaults to 'Feature Importance'.
        save_path (str | None, optional): The file path to save the plot. If `None`, the
            plot is displayed interactively. Defaults to `None`.
        **kwargs: Additional keyword arguments that can be passed to the underlying
            matplotlib functions, although none are currently used in the function body.

    Returns:
        None: The function does not return any value; it either displays the plot
            or saves it to a file.
    """
    # Convert tensor to numpy array for easier manipulation
    if isinstance(explanation, torch.Tensor):
        explanation = explanation.detach().cpu().numpy().flatten()

    # Create a figure with appropriate size
    plt.figure(figsize=(10, 6))

    # Sort features by their explanations
    indices = np.argsort(explanation)
    explanation = explanation[indices]
    if feature_names is not None:
        feature_names = np.array(feature_names)[indices]
    else:
        feature_names = indices

    # Plot vertical bars
    bars = plt.barh(
        range(len(explanation)), explanation, align="center", color="skyblue"
    )

    # Add feature names as y labels
    plt.yticks(range(len(explanation)), feature_names)

    # Add rounded explanation values on the bars
    for bar, expl in zip(bars, explanation):
        plt.text(
            bar.get_width() + 0.01,
            bar.get_y() + bar.get_height() / 2,
            f"{expl:.2f}",
            va="center",
        )

    # Add labels and title
    plt.xlabel("Explanation")
    plt.ylabel("Feature")
    plt.title(figure_title)

    # Save or show plot
    if save_path:
        plt.savefig(save_path)
    else:
        plt.show()
    plt.close()


def plot_explanation_multichannel_time_series(
    example: torch.Tensor | np.ndarray,
    explanation: torch.Tensor | np.ndarray,
    figure_title: str = "Heatmap",
    save_path: str | None = None,
    **kwargs,
):
    """Visualizes a multi-channel time series alongside a continuous heatmap.

    This function generates a figure with two plots for each channel of the time series.
    The left plot displays the raw time series data, while the right plot shows the
    same time series with a continuous color gradient applied to the line, where the
    color intensity corresponds to the explanation (heatmap) values. This helps in
    understanding the importance of different time points in the series.

    Args:
        example (torch.Tensor | np.ndarray): The input time series data. Expected
            shape is (batch_size, n_channels, n_timestamps), where batch_size is
            assumed to be 1.
        explanation (torch.Tensor | np.ndarray): The heatmap values corresponding to the
            time series. Expected to have the same shape as `example`.
        heatmap_title (str, optional): The title for the heatmap subplot.
            Defaults to 'Heatmap'.
        save_path (str | None, optional): The file path to save the plot. If `None`,
            the plot is displayed interactively. Defaults to `None`.
        **kwargs: Additional keyword arguments to be passed to the underlying
            matplotlib functions, although none are currently used in the function body.

    Returns:
        None: The function either displays the plot or saves it to a file.
    """
    if isinstance(example, torch.Tensor):
        example = example.detach().cpu().numpy()
    if isinstance(explanation, torch.Tensor):
        explanation = explanation.detach().cpu().numpy()
    # Extract the dimensions
    _, n_channels, n_timestamps = example.shape
    fig, axes = plt.subplots(n_channels, 2, figsize=(15, 5 * n_channels))

    # Handle the case where there is only one channel
    if n_channels == 1:
        axes = np.expand_dims(axes, axis=0)  # Convert from (2,) to (1, 2)

    for i in range(n_channels):
        # Extract time series and heatmap for the current channel
        ts = example[0, i]
        hm = explanation[0, i]

        # Plot the time series
        axes[i, 0].plot(np.arange(n_timestamps), ts, "o-")
        axes[i, 0].set_xlabel("Time", fontsize=14)
        axes[i, 0].set_title(f"Channel {i + 1} Time Series", fontsize=16)

        # Create a continuous line with color based on the heatmap
        for j in range(1, n_timestamps):
            # Color the segment of the line based on heatmap values
            segment = np.array([ts[j - 1], ts[j]])
            segment_times = np.array([j - 1, j])
            color = plt.cm.jet(hm[j])  # Using heatmap value directly
            axes[i, 1].plot(segment_times, segment, color=color, lw=2)

        # Set up the colorbar
        norm = plt.Normalize(vmin=0, vmax=1)  # Normalisation entre 0 et 1
        sm = plt.cm.ScalarMappable(cmap="jet", norm=norm)
        sm.set_array([])
        # cbar = fig.colorbar(sm, ax=axes[i, 1], label="Importance (X)")

        axes[i, 1].set_xlabel("Time", fontsize=14)
        axes[i, 1].set_title(f"{figure_title} (Channel {i + 1})", fontsize=16)

    plt.tight_layout()
    # Save or show plot
    if save_path:
        plt.savefig(save_path)
    else:
        plt.show()
    plt.close()


def plot_explanation_comparison_image(
    explanations: dict[int, dict[str, np.ndarray]],
    targets: dict[int, str],
    save_path: str | Path,
) -> None:
    """Plots a comparison image of different explanations (heatmaps) for multiple examples.

    Each row in the resulting figure represents a single input example (image). The
    first column is the original **Input** image, and subsequent columns display
    different explanation heatmaps overlaid on the input image.

    Args:
        explanations (dict[int, dict[str, np.ndarray]]): A dictionary where keys are
            example indices (int) and values are dictionaries containing the
            explanation components. The inner dictionary must contain the key **"Input"**
            with the original image data (numpy array), and other keys
            representing different explainers (e.g., "First_explainer") with their
            corresponding heatmap data (numpy array). All arrays are expected
            to be image-like (e.g. 3D).
            Example structure:
            `{0: {"Input": np.ndarray, "ExplainerA": np.ndarray, ...}, 1: ...}`
        targets (dict[int, str]): a dict that maps image idx with its corresponding class label
        save_path (str | Path): The file path where the comparison figure will be saved (e.g., 'comparison.png').
    """
    column_names = ["Label"] + sorted(
        explanations[0].keys(), key=lambda x: (x != "Input", x)
    )
    fig, axes = plt.subplots(
        len(explanations), len(column_names), figsize=(10, 10)
    )
    for col_idx, column in enumerate(column_names):
        if col_idx > 0:
            axes[0, col_idx].set_title(column, fontsize=12, fontweight="bold")
        for row_idx, elements in explanations.items():
            if col_idx == 0:
                axes[row_idx, col_idx].text(
                    x=0.5,
                    y=0.5,
                    s=targets.get(row_idx, ""),
                    fontsize=12,
                    ha="center",
                    va="center",
                )
                axes[row_idx, col_idx].axis("off")
            else:
                axes[row_idx, col_idx].imshow(elements["Input"])
                axes[row_idx, col_idx].imshow(
                    elements[column], cmap="jet", alpha=0.5
                )
                axes[row_idx, col_idx].axis("off")
    plt.tight_layout()
    fig.savefig(save_path)
