"""Metrics visualization plots."""

import numpy as np
import plotly.graph_objects as go
import xarray as xr


def plot_confusion_matrix(cm_data: xr.DataArray, title: str = "Confusion Matrix", normalize: str = "none") -> go.Figure:
    """Plot an interactive confusion matrix heatmap.

    Args:
        cm_data: XArray DataArray with confusion matrix data (dimensions: true_label, predicted_label).
        title: Title for the plot.
        normalize: Normalization mode - "none", "true", or "pred".

    Returns:
        Plotly figure with the interactive confusion matrix heatmap.

    """
    # Get the data as numpy array
    cm_array = cm_data.values.astype(float)
    labels = cm_data.coords["true_label"].values.tolist()

    # Store original counts for display
    cm_counts = cm_data.values

    # Apply normalization
    if normalize == "true":
        # Normalize over true labels (rows) - each row sums to 1
        row_sums = cm_array.sum(axis=1, keepdims=True)
        cm_normalized = np.divide(cm_array, row_sums, where=row_sums != 0)
        colorbar_title = "Proportion"
    elif normalize == "pred":
        # Normalize over predicted labels (columns) - each column sums to 1
        col_sums = cm_array.sum(axis=0, keepdims=True)
        cm_normalized = np.divide(cm_array, col_sums, where=col_sums != 0)
        colorbar_title = "Proportion"
    else:
        # No normalization
        cm_normalized = cm_array
        colorbar_title = "Count"

    # Create annotations for the heatmap
    annotations = []
    for i, true_label in enumerate(labels):
        for j, pred_label in enumerate(labels):
            count = int(cm_counts[i, j])
            normalized_val = cm_normalized[i, j]

            # Format text based on normalization mode
            if normalize == "none":
                # Show count and percentage of total
                total = cm_counts.sum()
                pct = (count / total * 100) if total > 0 else 0
                text = f"{count}<br>({pct:.1f}%)"
            else:
                # Show percentage only for normalized versions
                text = f"{normalized_val:.1%}"

            # Determine text color based on normalized value
            threshold = cm_normalized.max() / 2 if cm_normalized.max() > 0 else 0.5
            text_color = "white" if normalized_val > threshold else "black"

            annotations.append(
                {
                    "x": pred_label,
                    "y": true_label,
                    "text": text,
                    "showarrow": False,
                    "font": {"size": 10, "color": text_color},
                }
            )

    # Create the heatmap with normalized values for coloring
    fig = go.Figure(
        data=go.Heatmap(
            z=cm_normalized,
            x=labels,
            y=labels,
            colorscale="Blues",
            colorbar={"title": colorbar_title},
            hoverongaps=False,
            hovertemplate="True: %{y}<br>Predicted: %{x}<br>Count: %{customdata}<extra></extra>",
            customdata=cm_counts,
        )
    )

    # Add annotations
    fig.update_layout(
        annotations=annotations,
        xaxis={"title": "Predicted Label", "side": "bottom"},
        yaxis={"title": "True Label", "autorange": "reversed"},
        width=600,
        height=550,
    )

    return fig
