"""Produce reproducible figures and a paper comparison from saved test predictions."""
from __future__ import annotations

import csv
import json
import shutil
from pathlib import Path

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

from common import ROOT, arguments, configuration, output_path


def main():
    parser = arguments(__doc__)
    parser.add_argument("--sample", type=int, default=19)
    args = parser.parse_args()
    cfg = configuration(args)
    directory = output_path(cfg["output"]["results_dir"])
    metrics = json.loads((directory / "metrics.json").read_text())
    summary = json.loads((directory / "training_summary.json").read_text())
    with (directory / "history.csv").open() as stream:
        history = list(csv.DictReader(stream))
    if len(history) != 500 or [int(row["epoch"]) for row in history] != list(range(1, 501)):
        raise ValueError("Results require one contiguous full 500-epoch training history")
    data = torch.load(directory / "test_predictions.pt", map_location="cpu", weights_only=True)
    if len(data["prediction"]) != 100:
        raise ValueError("Expected all 100 test predictions")
    epochs = np.array([int(row["epoch"]) for row in history])
    energy = np.array([float(row["train_energy"]) for row in history])
    error = np.array([float(row["test_relative_l2"]) for row in history]) * 100
    fig, axes = plt.subplots(1, 2, figsize=(11, 4), constrained_layout=True)
    axes[0].plot(epochs, energy, label="Train variational energy")
    axes[0].plot(epochs, [float(row["test_energy"]) for row in history], label="Test variational energy", alpha=0.8)
    axes[0].set(xlabel="Epoch", ylabel="Energy", title="Unlabeled training objective")
    axes[0].legend()
    axes[1].semilogy(epochs, error, label="Test relative L2")
    axes[1].axhline(cfg["reference"]["table1_mean_percent"], color="black", linestyle="--", label="Paper Table 1 mean")
    axes[1].scatter([metrics["checkpoint_epoch"]], [metrics["relative_l2_mean_percent"]], color="red", s=30, label="Selected checkpoint", zorder=3)
    axes[1].set(xlabel="Epoch", ylabel="Mean relative L2 (%)", title="100 fixed test cases")
    axes[1].legend()
    for extension in ("png", "pdf"):
        fig.savefig(directory / f"training_curves.{extension}", dpi=180)
    plt.close(fig)
    errors = data["relative_l2"].numpy() * 100
    fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)
    axes[0].hist(errors, bins=16, color="#2879a5", edgecolor="white")
    axes[0].axvline(errors.mean(), color="red", label=f"Mean {errors.mean():.3f}%")
    axes[0].set(xlabel="Relative L2 (%)", ylabel="Test cases", title="Test error distribution")
    axes[0].legend()
    axes[1].boxplot(errors, labels=["VINO reproduction"])
    axes[1].axhline(cfg["reference"]["table1_mean_percent"], linestyle="--", color="black", label="Paper mean")
    axes[1].set(ylabel="Relative L2 (%)")
    axes[1].legend()
    fig.savefig(directory / "test_error_distribution.png", dpi=180)
    plt.close(fig)
    order = np.argsort(errors)
    selections = list(dict.fromkeys([args.sample, int(order[0]), int(order[len(order)//2]), int(order[-1])]))
    for index in selections:
        if index < 0 or index >= 100:
            raise ValueError("Sample index must be in [0,99]")
        coefficient = data["nodal_permeability"][index, ..., 0].numpy()
        reference = data["reference"][index, ..., 0].numpy()
        predicted = data["prediction"][index, ..., 0].numpy()
        low, high = min(reference.min(), predicted.min()), max(reference.max(), predicted.max())
        fig, axes = plt.subplots(1, 4, figsize=(14, 3.7), constrained_layout=True)
        for ax, value, title, cmap, limits in zip(axes,
                (coefficient, reference, predicted, np.abs(predicted-reference)),
                ("Nodal permeability", "Reference head", "VINO head", "Absolute error"),
                ("viridis", "viridis", "viridis", "magma"),
                ((None, None), (low, high), (low, high), (0, None))):
            plot = ax.imshow(value.T, origin="lower", extent=[0,1,0,1], cmap=cmap, vmin=limits[0], vmax=limits[1])
            ax.set(xlabel="x", ylabel="y", title=title)
            fig.colorbar(plot, ax=ax, shrink=0.8)
        fig.suptitle(f"Original test case {index}: relative L2 = {errors[index]:.4f}%")
        fig.savefig(directory / f"test_case_{index:03d}.png", dpi=180)
        plt.close(fig)
    report = f"""# VINO Darcy Flow experiment

Completed **500 epochs**, using the original **1000 training / 100 test** samples at **64×64** resolution.

| Measurement | Reproduction | Paper Table 1 |
|---|---:|---:|
| Test relative L2 mean | {metrics['relative_l2_mean_percent']:.6f}% | 0.93% |
| Test relative L2 sample standard deviation | {metrics['relative_l2_std_percent']:.6f}% | 0.31% |
| Training wall time | {summary['training_seconds']:.2f} s | 244.2 s (A100) |

Selected checkpoint: epoch **{metrics['checkpoint_epoch']}**, `weight/best_model.pt`.
Final epoch 500 mean relative L2: **{summary['last_test_relative_l2']*100:.6f}%**.
Zero Dirichlet boundary maximum absolute value: **{metrics['boundary_max_abs']}**.
Mean error difference from paper: **{metrics['relative_l2_mean_percent']-0.93:+.6f} percentage points**.

Training uses only permeability and the variational energy. U_train never enters optimization.
The test labels select the best checkpoint as explicitly requested; the reported best-checkpoint test score therefore includes model selection on that same test set. The epoch-500 score is also provided.
The standard deviation is across 100 test samples, not across independent training seeds.

Paper: https://arxiv.org/html/2411.06587 (Section 2.1.3, Table 1, Section 4.2).
Implementation is independent. Official source was viewed online only for paper omissions; no repository code was downloaded or copied.
Exact Q1 rectangular-element integration follows the paper rather than the ambiguous manual Darcy expansion in the official helper. The coefficient is bilinearly interpolated within each element after the official missing-detail nodal averaging step.
Network and optimizer details, assumptions, and source conflicts are recorded in `reproduction_spec.md`.
The paper describes a permeability contrast of 4, whereas the required author file contains the two values 4 and 12 (contrast 3). The original data were used without rescaling or replacement, as requested.
Hardware/software are in `environment.json`; timing includes per-epoch testing and checkpoint I/O and is not a direct hardware-controlled comparison.

Dataset SHA256: `{metrics['data_sha256']}`.
Files: `history.csv`, `metrics.json`, `per_sample_metrics.csv`, `test_predictions.pt`, `training_curves.png`, `test_error_distribution.png`, and selected `test_case_*.png`.
"""
    (directory / "report.md").write_text(report)
    paper_spec = ROOT.parents[2] / ".paper2code_work/2411.06587/reproduction_spec.md"
    if paper_spec.exists():
        shutil.copyfile(paper_spec, directory / "reproduction_spec.md")
    print(report, flush=True)


if __name__ == "__main__":
    main()