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 load_config, resolve_output_path, save_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml"))
    parser.add_argument("--predictions", default=None)
    parser.add_argument("--num-samples", type=int, default=4)
    return parser.parse_args()


def rel_l2(pred: np.ndarray, target: np.ndarray) -> np.ndarray:
    diff = np.linalg.norm((pred - target).reshape(pred.shape[0], -1), axis=1)
    denom = np.linalg.norm(target.reshape(target.shape[0], -1), axis=1)
    return diff / np.maximum(denom, 1e-12)


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"])
    pred_path = Path(args.predictions) if args.predictions else results_dir / "test_predictions.npz"
    data = np.load(pred_path)
    inputs = data["inputs"]
    preds = data["predictions"]
    targets = data["targets"]
    errors = rel_l2(preds, targets) * 100.0

    sample_count = min(args.num_samples, preds.shape[0])
    fig, axes = plt.subplots(sample_count, 4, figsize=(13, 3.1 * sample_count), constrained_layout=True)
    if sample_count == 1:
        axes = axes[None, :]
    for i in range(sample_count):
        panels = [
            (inputs[i, 0], "input t=0"),
            (targets[i, 0], "target t=5"),
            (preds[i, 0], "prediction"),
            (np.abs(preds[i, 0] - targets[i, 0]), f"abs error rel={errors[i]:.2f}%"),
        ]
        for ax, (arr, title) in zip(axes[i], panels):
            im = ax.imshow(arr, origin="lower", cmap="RdBu_r" if title != "abs error" else "magma")
            ax.set_title(title)
            ax.set_xticks([])
            ax.set_yticks([])
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
    fig_path = results_dir / "prediction_examples.png"
    fig.savefig(fig_path, dpi=180)
    plt.close(fig)

    hist_fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
    ax.hist(errors, bins=24, color="#386cb0", edgecolor="white")
    ax.axvline(errors.mean(), color="#d95f02", linewidth=2, label=f"mean={errors.mean():.2f}%")
    ax.axvline(cfg["experiment"]["paper_reference_test_percent"], color="#1b9e77", linewidth=2, linestyle="--", label="paper CROP 4.06%")
    ax.set_xlabel("relative L2 (%)")
    ax.set_ylabel("test samples")
    ax.legend()
    hist_path = results_dir / "relative_l2_histogram.png"
    hist_fig.savefig(hist_path, dpi=180)
    plt.close(hist_fig)

    save_json(
        results_dir / "result_summary.json",
        {
            "status": "complete",
            "prediction_file": str(pred_path),
            "mean_relative_l2_percent": float(errors.mean()),
            "std_relative_l2_percent": float(errors.std(ddof=0)),
            "min_relative_l2_percent": float(errors.min()),
            "max_relative_l2_percent": float(errors.max()),
            "figures": [str(fig_path), str(hist_path)],
            "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"],
            "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"],
        },
    )
    print(
        f"result_complete mean_relative_l2={errors.mean():.6f}% "
        f"std={errors.std(ddof=0):.6f}% figures={[str(fig_path), str(hist_path)]}",
        flush=True,
    )


if __name__ == "__main__":
    main()