"""
Visualization utilities for 2D numerical simulations.
Provides functions to generate and save side-by-side plots of model predictions
against ground truth, with options for custom titles, labels, color maps, and
batch/video frame sequencing.

Functions:
    plot_NS_2D_trainer: Generate prediction plots during training and save to disk.
    plot_NS_2D_multiple: Create a figure with multiple named data plots.
    plot_NS_2D: Plot a single 2D field on an Axes with optional metadata.
"""

import matplotlib.pyplot as plt
from utils import is_binary
import torch
import os


def plot_NS_2D_trainer(
    model, plot_input, plot_target, plot_path, step, freq_plot, show, cmap
):
    """
    Generate and save a plot of model predictions versus targets at training intervals.
    """
    with torch.no_grad():
        x, y_pred, y = (
            plot_input[:, :, -1],
            model(plot_input.unsqueeze(0)).squeeze(0),
            plot_target,
        )

        plot = plot_NS_2D_multiple(show=show, x=x, y_pred=y_pred, y=y, cmap=cmap)

        os.makedirs(plot_path, exist_ok=True)
        path_new_image = os.path.join(plot_path, f"step_{step // freq_plot:05d}.png")
        plot.savefig(path_new_image)
        plt.close()


def plot_NS_2D_multiple(
    title=None,
    x_label=None,
    y_label=None,
    show=False,
    idx=None,
    cmap=None,
    **named_data,
):
    """
    Create a multi-panel figure of 2D fields.
    """
    n_plots = sum(data is not None for data in named_data.values())
    if n_plots == 0:
        return plt

    fig, axes = plt.subplots(1, n_plots, figsize=(6 * n_plots, 6))  # 1 row, N columns
    if n_plots == 1:
        axes = [axes]  # Make it iterable if only one plot

    for ax, (name, data) in zip(axes, named_data.items()):
        if data is not None:
            plot_NS_2D(
                data, idx, label=name, ax=ax, cmap=cmap
            )  # (grid, data, idx, label=name)
            ax.set_title(name)
            ax.set_xlabel(x_label if x_label is not None else "")
            ax.set_ylabel(y_label if y_label is not None else "")
            ax.legend()

    fig.suptitle(
        title if title is not None else f"Darcy plots for data point {idx}", y=0.95
    )
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def plot_NS_2D(u, idx=None, label=None, ax=None, cmap=None):
    """
    Plot a single 2D field on a given axis.
    """
    if is_binary(u[:, 0]):
        cmap = "gray"
    else:
        cmap = cmap

    u_plt = u.detach().cpu()
    if idx is not None:
        u_plt = u_plt[idx]
    ax.imshow(u_plt, cmap=cmap)
    if label:
        ax.set_title(label)
