from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from common import ROOT, load_config, resolve, save_json


def plot_case(xy: np.ndarray, pred: np.ndarray, target: np.ndarray, out_path: Path, title: str) -> None:
    x = xy[..., 0]
    y = xy[..., 1]
    p = pred[..., 0]
    t = target[..., 0]
    err = p - t
    vmin = float(min(p.min(), t.min()))
    vmax = float(max(p.max(), t.max()))
    fig, axes = plt.subplots(1, 3, figsize=(13, 3.6), constrained_layout=True)
    for ax, field, name, lo, hi, cmap in [
        (axes[0], t, "target", vmin, vmax, "coolwarm"),
        (axes[1], p, "prediction", vmin, vmax, "coolwarm"),
        (axes[2], err, "error", float(-np.max(np.abs(err))), float(np.max(np.abs(err))), "coolwarm"),
    ]:
        mesh = ax.pcolormesh(x, y, field, shading="auto", cmap=cmap, vmin=lo, vmax=hi)
        ax.set_aspect("equal", adjustable="box")
        ax.set_title(name)
        ax.set_xticks([])
        ax.set_yticks([])
        fig.colorbar(mesh, ax=ax, fraction=0.046, pad=0.02)
    fig.suptitle(title)
    fig.savefig(out_path, dpi=180)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(description="Create HO-FNO Pipe visualizations from inference arrays.")
    parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml"))
    parser.add_argument("--num-cases", type=int, default=8)
    args = parser.parse_args()

    cfg = load_config(args.config)
    output_dir = resolve(cfg["paths"]["output_dir"])
    results_dir = resolve(cfg["paths"]["results_dir"], output_dir)
    pred_path = results_dir / "predictions.npy"
    target_path = results_dir / "targets.npy"
    input_path = results_dir / "inputs.npy"
    rel_path = results_dir / "relative_l2_per_sample.npy"
    for path in [pred_path, target_path, input_path, rel_path]:
        if not path.exists():
            raise FileNotFoundError(f"missing required inference artifact: {path}")

    pred = np.load(pred_path)
    target = np.load(target_path)
    inputs = np.load(input_path)
    rel = np.load(rel_path)
    fig_dir = results_dir / "figures"
    fig_dir.mkdir(parents=True, exist_ok=True)
    order = np.argsort(rel)
    chosen = []
    if len(order) > 0:
        chosen.extend([int(order[0]), int(order[len(order) // 2]), int(order[-1])])
    for idx in range(min(args.num_cases, len(rel))):
        if idx not in chosen:
            chosen.append(idx)
    chosen = chosen[: max(1, args.num_cases)]
    figures = []
    for idx in chosen:
        out_path = fig_dir / f"pipe_case_{idx:03d}.png"
        plot_case(inputs[idx], pred[idx], target[idx], out_path, f"Pipe case {idx} | rel L2={rel[idx]:.6f}")
        figures.append(str(out_path))

    summary = {
        "figure_count": len(figures),
        "figures": figures,
        "mean_relative_l2": float(rel.mean()),
        "best_case": int(order[0]) if len(order) else None,
        "median_case": int(order[len(order) // 2]) if len(order) else None,
        "worst_case": int(order[-1]) if len(order) else None,
    }
    save_json(results_dir / "result_summary.json", summary)
    print(f"result_complete figures={len(figures)} summary={results_dir / 'result_summary.json'}", flush=True)


if __name__ == "__main__":
    main()