from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import matplotlib

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_config(path: str | Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def resolve(path: str | Path) -> Path:
    path = Path(path)
    return path if path.is_absolute() else PROJECT_ROOT / path


def load_json(path: Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def save_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


def plot_sample(sample_id: int, inputs: np.ndarray, targets: np.ndarray, preds: np.ndarray, channel_names: list[str], out_dir: Path) -> Path:
    fig, axes = plt.subplots(4, len(channel_names), figsize=(3.2 * len(channel_names), 10.0), constrained_layout=True)
    rows = [("Boussinesq input", inputs), ("Compressible target", targets), ("FNO prediction", preds), ("Abs error", np.abs(preds - targets))]
    for row, (title, array) in enumerate(rows):
        for col, channel in enumerate(channel_names):
            ax = axes[row, col]
            image = array[sample_id, col]
            vmax = 1.0 if row < 3 else max(float(image.max()), 1.0e-8)
            im = ax.imshow(image, origin="lower", cmap="viridis" if row < 3 else "magma", vmin=0.0, vmax=vmax)
            ax.set_xticks([])
            ax.set_yticks([])
            if row == 0:
                ax.set_title(channel)
            if col == 0:
                ax.set_ylabel(title)
            if row == 3:
                mse = float(np.mean((preds[sample_id, col] - targets[sample_id, col]) ** 2))
                ax.set_xlabel(f"MSE={mse:.2e}")
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / f"sample_{sample_id:03d}.png"
    fig.savefig(path, dpi=160)
    plt.close(fig)
    return path


def main() -> None:
    parser = argparse.ArgumentParser(description="Summarize metrics and render NeuralConvection FNO figures.")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml"))
    args = parser.parse_args()

    config = load_config(args.config)
    viz_cfg = config["visualization"]
    predictions = np.load(resolve(viz_cfg["predictions_path"]))
    preds = predictions["pred_norm"]
    targets = predictions["target_norm"]
    inputs = predictions["input_norm"]
    channel_names = [str(x) for x in predictions["channel_names"].tolist()]
    metrics = load_json(resolve(viz_cfg["metrics_path"]))
    out_dir = resolve(viz_cfg["output_dir"])
    figures = []
    for sample_id in range(min(int(viz_cfg["max_samples"]), preds.shape[0])):
        figures.append(str(plot_sample(sample_id, inputs, targets, preds, channel_names, out_dir)))

    comparison = {}
    paper_channels = config["paper_table5_compact_fno_2d"]["channels"]
    for name in channel_names:
        current = metrics["per_channel"][name]
        paper = paper_channels[name]
        comparison[name] = {
            "mse": current["mse"],
            "paper_mse": paper["mse"],
            "mse_ratio_to_paper": current["mse"] / paper["mse"],
            "ssim": current["ssim"],
            "paper_ssim": paper["ssim"],
            "ssim_delta": current["ssim"] - paper["ssim"],
        }
    summary = {
        "status": "success",
        "predictions_path": str(resolve(viz_cfg["predictions_path"])),
        "metrics_path": str(resolve(viz_cfg["metrics_path"])),
        "figures": figures,
        "table5_comparison": comparison,
    }
    save_json(resolve(viz_cfg["summary_path"]), summary)
    print(f"result_summary={resolve(viz_cfg['summary_path'])}", flush=True)
    for channel, values in comparison.items():
        print(
            f"{channel}: mse={values['mse']:.8e} paper={values['paper_mse']:.8e} "
            f"ratio={values['mse_ratio_to_paper']:.3f} ssim={values['ssim']:.6f} paper_ssim={values['paper_ssim']:.6f}",
            flush=True,
        )
    print(f"figures={len(figures)} saved_to={out_dir}", flush=True)


if __name__ == "__main__":
    main()