"""Create diagnostic plots from Samudra NPZ rollout outputs."""

try:
    from ._bootstrap import ROOT
except ImportError:
    from _bootstrap import ROOT

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

DEPTHS = (2.5, 10.0, 22.5, 40.0, 65.0, 105.0, 165.0, 250.0, 375.0, 550.0, 775.0, 1050.0, 1400.0, 1850.0, 2400.0, 3100.0, 4000.0, 5000.0, 6000.0)
PROGNOSTIC_CHANNELS = tuple(
    f"{name}_{level}"
    for name in ("uo", "vo", "thetao", "so")
    for level in range(len(DEPTHS))
) + ("zos",)


def plot_rollout(prediction_path: str, output_dir: str) -> None:
    with np.load(prediction_path) as data:
        predictions = np.asarray(data["predictions"])
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
    axes[0].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("zos")], cmap="coolwarm", aspect="auto")
    axes[0].set_title("SSH forecast")
    axes[1].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("thetao_0")], cmap="turbo", aspect="auto")
    axes[1].set_title("Surface potential temperature")
    fig.savefig(output / "forecast_maps.png", dpi=150)
    plt.close(fig)
    indices = [PROGNOSTIC_CHANNELS.index(f"thetao_{i}") for i in range(len(DEPTHS))]
    profile = predictions[:, indices].mean(axis=(0, 2, 3))
    fig, ax = plt.subplots(figsize=(5, 5), constrained_layout=True)
    ax.plot(profile, DEPTHS, marker="o")
    ax.invert_yaxis()
    ax.set(xlabel="potential temperature", ylabel="depth (m)", title="Mean temperature profile")
    fig.savefig(output / "temperature_profile.png", dpi=150)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--prediction", default="./result/output/prediction.npz")
    parser.add_argument("--output-dir", default="./result")
    args = parser.parse_args()
    plot_rollout(args.prediction, args.output_dir)


if __name__ == "__main__":
    main()
    print("save to ./result/")