"""Render a result's plot series to a PNG."""

from __future__ import annotations

from pathlib import Path

import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

from zenith.engines.base import Result  # noqa: E402


def save_plot(result: Result, path: Path) -> None:
    items = list(result.series.items())
    if not items:
        raise ValueError(f"{result.sim}/{result.op} has no plot series")
    fig, axes = plt.subplots(1, len(items), figsize=(5 * len(items), 4), squeeze=False)
    for ax, (name, s) in zip(axes[0], items):
        if "z" in s:  # map grid: rows run north to south
            z = np.array([[np.nan if v is None else v for v in row] for row in s["z"]], dtype=float)
            im = ax.imshow(z, cmap="magma", extent=(s["x"][0], s["x"][-1], s["y"][-1], s["y"][0]), aspect="auto")
            fig.colorbar(im, ax=ax, label=s["z_label"])
        else:
            ax.plot(s["x"], s["y"], marker="o" if len(s["x"]) == 1 else None)
        ax.set_title(name.replace("_", " "))
        ax.set_xlabel(s["x_label"])
        ax.set_ylabel(s["y_label"])
        if name.startswith("orbit_xy"):
            ax.set_aspect("equal")
            ax.plot([0], [0], "*", color="orange", markersize=12)
    fig.suptitle(f"{result.sim} / {result.op}")
    fig.tight_layout()
    fig.savefig(path, dpi=100)
    plt.close(fig)