#!/usr/bin/env python3
"""Create the final FAST-DeepONet Darcy figures and reproduction report."""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import matplotlib

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def load_config(path: Path) -> dict[str, object]:
    """Load visualization configuration without importing the Torch runtime."""
    with path.resolve().open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    if not isinstance(config, dict):
        raise ValueError(f"Config must be a mapping: {path}")
    return config


def atomic_save_figure(figure: plt.Figure, path: Path, dpi: int = 180) -> None:
    temporary = path.with_name(path.stem + ".tmp" + path.suffix)
    figure.savefig(temporary, dpi=dpi, bbox_inches="tight")
    os.replace(temporary, path)


def plot_training(history: list[dict[str, float]], results_dir: Path) -> None:
    epochs = np.asarray([row["epoch"] for row in history])
    train_mse = np.asarray([row["train_mse"] for row in history])
    tune_mse = np.asarray([row["tune_mse"] for row in history])
    train_total = np.asarray([row["train_total_loss"] for row in history])
    figure, axis = plt.subplots(figsize=(7.5, 4.8))
    axis.semilogy(epochs, train_total, label="train objective", linewidth=1.4)
    axis.semilogy(epochs, train_mse, label="train MSE", linewidth=1.4)
    axis.semilogy(epochs, tune_mse, label="tune MSE", linewidth=1.6)
    axis.set_xlabel("Epoch")
    axis.set_ylabel("Loss")
    axis.set_title("FAST-DeepONet Darcy training")
    axis.grid(True, which="both", alpha=0.25)
    axis.legend()
    atomic_save_figure(figure, results_dir / "training_curves.png")
    plt.close(figure)


def plot_sample_metrics(metrics: dict[str, object], results_dir: Path) -> None:
    values = np.asarray([sample["relative_l2"] for sample in metrics["samples"]], dtype=float)
    figure, axis = plt.subplots(figsize=(7.5, 4.5))
    bars = axis.bar(np.arange(values.size), values, color="#4C78A8")
    axis.axhline(float(metrics["mean_relative_l2"]), color="#E45756", linestyle="--", label="mean")
    axis.set_xticks(np.arange(values.size), [str(index) for index in range(values.size)])
    axis.set_xlabel("Test sample index")
    axis.set_ylabel("Relative L2")
    axis.set_title("Per-sample test error")
    axis.grid(True, axis="y", alpha=0.25)
    axis.legend()
    for bar, value in zip(bars, values):
        axis.text(bar.get_x() + bar.get_width() / 2, value, f"{value:.3f}", ha="center", va="bottom", fontsize=8)
    atomic_save_figure(figure, results_dir / "test_relative_l2.png")
    plt.close(figure)


def plot_field_comparison(prediction_path: Path, grid_size: int, results_dir: Path) -> int:
    with np.load(prediction_path, allow_pickle=False) as data:
        errors = np.asarray(data["relative_l2"])
        sample_index = int(np.argsort(errors)[errors.size // 2])
        log_permeability = np.asarray(data["inputs"][sample_index]).reshape(grid_size, grid_size)
        target = np.asarray(data["targets"][sample_index]).reshape(grid_size, grid_size)
        prediction = np.asarray(data["predictions"][sample_index]).reshape(grid_size, grid_size)
        mask = np.asarray(data["masks"][sample_index]).reshape(grid_size, grid_size)
        filename = str(data["filenames"][sample_index])
    absolute_error = np.abs(target - prediction)
    figure, axes = plt.subplots(1, 4, figsize=(16, 3.8), constrained_layout=True)
    panels = [
        (log_permeability, "log permeability", "coolwarm"),
        (target, "pressure target", "coolwarm"),
        (prediction, "pressure prediction", "coolwarm"),
        (absolute_error, "absolute error", "magma"),
    ]
    pressure_min = min(float(target.min()), float(prediction.min()))
    pressure_max = max(float(target.max()), float(prediction.max()))
    for panel_index, (axis, (field, title, cmap)) in enumerate(zip(axes, panels)):
        kwargs = {"origin": "lower", "extent": (0, 1, 0, 1), "cmap": cmap}
        if panel_index in (1, 2):
            kwargs.update(vmin=pressure_min, vmax=pressure_max)
        image = axis.imshow(field.T, **kwargs)
        if np.any(mask > 0.5) and np.any(mask <= 0.5):
            axis.contour(
                np.linspace(0, 1, grid_size),
                np.linspace(0, 1, grid_size),
                mask.T,
                levels=[0.5],
                colors="white" if panel_index != 0 else "black",
                linewidths=0.5,
                alpha=0.7,
            )
        axis.set_title(title)
        axis.set_xlabel("x")
        if panel_index == 0:
            axis.set_ylabel("y")
        figure.colorbar(image, ax=axis, shrink=0.82)
    figure.suptitle(f"Test sample {sample_index}: {filename} (relative L2={errors[sample_index]:.5f})")
    atomic_save_figure(figure, results_dir / "darcy_field_comparison.png")
    plt.close(figure)
    return sample_index


def write_report(
    path: Path,
    training_summary: dict[str, object],
    metrics: dict[str, object],
    representative_index: int,
) -> None:
    lines = [
        "# FAST-DeepONet Darcy reproduction report",
        "",
        "## Formal training",
        "",
        f"- Status: `{training_summary['status']}`",
        f"- Completed epochs: {training_summary['completed_epochs']} / {training_summary['requested_epochs']}",
        f"- Best epoch: {training_summary['best_epoch']}",
        f"- Best tune MSE: {float(training_summary['best_tune_mse']):.10e}",
        f"- Effective factor rank: {training_summary['effective_factor_rank']}",
        f"- Device: `{training_summary['device_name']}`",
        f"- SLURM job: `{training_summary['slurm_job_id']}`",
        f"- Checkpoint: `{training_summary['checkpoint']}`",
        "",
        "## Test relative L2",
        "",
    ]
    for sample in metrics["samples"]:
        lines.append(
            f"- {int(sample['sample_index']):02d} `{sample['filename']}`: {float(sample['relative_l2']):.10e}"
        )
    lines.extend(
        [
            "",
            f"Mean relative L2: **{float(metrics['mean_relative_l2']):.10e}**",
            "",
            "## Artifacts and interpretation",
            "",
            f"- The four-panel field figure uses representative sample index {representative_index}.",
            "- `training_curves.png` shows the actual 300-epoch objective/train MSE/tune MSE history.",
            "- `test_relative_l2.png` shows all ten fixed test fields.",
            "- Calibration data were not used because this task requests point prediction, not UQ calibration.",
            "- This run uses 10 test fields and one model seed. Table 2 reports 100 independent test fields and summarizes five fits, so its 0.0925±0.0019 value is not directly comparable to this fixed-data run.",
            "",
        ]
    )
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text("\n".join(lines), encoding="utf-8")
    os.replace(temporary, path)


def run(config: dict[str, object]) -> None:
    results_dir = Path(config["paths"]["results_dir"])
    history_path = results_dir / "training_history.json"
    summary_path = results_dir / "training_summary.json"
    metrics_path = results_dir / "test_metrics.json"
    predictions_path = results_dir / "test_predictions.npz"
    for path in (history_path, summary_path, metrics_path, predictions_path):
        if not path.is_file():
            raise FileNotFoundError(f"Required result is missing: {path}")
    history = json.loads(history_path.read_text(encoding="utf-8"))
    training_summary = json.loads(summary_path.read_text(encoding="utf-8"))
    metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
    if len(history) != 300 or int(training_summary["completed_epochs"]) != 300:
        raise ValueError("Visualization requires the completed 300-epoch formal run")
    if len(metrics["samples"]) != 10:
        raise ValueError("Visualization requires all ten test metrics")
    plot_training(history, results_dir)
    plot_sample_metrics(metrics, results_dir)
    representative_index = plot_field_comparison(
        predictions_path, int(config["data"]["grid_size"]), results_dir
    )
    write_report(results_dir / "reproduction_report.md", training_summary, metrics, representative_index)
    print(
        f"visualization_complete mean_relative_l2={float(metrics['mean_relative_l2']):.10e} "
        f"results_dir={results_dir}",
        flush=True,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    run(load_config(args.config))


if __name__ == "__main__":
    main()