"""Plot training curves and physical-space validation predictions."""

import argparse
import csv
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from common import load_config, package_path, save_json


def render(config):
    output = package_path(config["results"]["directory"])
    data = np.load(output/"predictions.npz")
    metrics = json.loads((output/"metrics.json").read_text())
    with open(output/"history.csv") as stream:
        rows = list(csv.DictReader(stream))
    epoch = np.array([int(r["epoch"]) for r in rows])
    fig, axes = plt.subplots(1, 2, figsize=(11, 4), constrained_layout=True)
    for name, label in [("train_loss","Training"),("val_loss","Validation")]:
        axes[0].semilogy(epoch,[float(r[name]) for r in rows],label=label)
    axes[0].set(xlabel="Epoch",ylabel="MSE + relative L2 (normalized)")
    axes[0].legend()
    axes[1].semilogy(epoch,[float(r["val_relative_l2"]) for r in rows],color="tab:blue",label="Validation")
    axes[1].axhline(0.0109,color="tab:red",linestyle="--",label="Paper Table 1: 0.0109")
    axes[1].set(xlabel="Epoch",ylabel="Relative L2 (physical space)")
    axes[1].legend()
    for ax in axes: ax.grid(alpha=0.2)
    fig.savefig(output/"training_curves.png",dpi=200)
    fig.savefig(output/"training_curves.pdf")
    plt.close(fig)
    errors=data["relative_l2"]
    chosen=[0,int(np.argsort(errors)[len(errors)//2]),int(np.argmax(errors))]
    fig, axes=plt.subplots(3,2,figsize=(12,10),constrained_layout=True)
    x=data["x"]
    for row,index in enumerate(chosen):
        initial=data["initial"][index,:,0]; target=data["target"][index,:,0]; pred=data["prediction"][index,:,0]
        axes[row,0].plot(x,initial,color="0.6",linewidth=1,label="Initial u0")
        axes[row,0].plot(x,target,color="black",linewidth=1.5,label="Reference u(T)")
        axes[row,0].plot(x,pred,color="tab:blue",linestyle="--",label="LiNO")
        axes[row,0].set_title(f"Source sample {data['sample_indices'][index]}, relative L2={errors[index]:.4g}")
        axes[row,0].set(xlabel="x",ylabel="u (dataset units)")
        axes[row,0].legend()
        axes[row,1].plot(x,pred-target,color="tab:red")
        axes[row,1].axhline(0,color="0.5",linewidth=0.5)
        axes[row,1].set(xlabel="x",ylabel="Prediction - reference")
    fig.savefig(output/"validation_examples.png",dpi=200)
    fig.savefig(output/"validation_examples.pdf")
    plt.close(fig)
    fig, ax=plt.subplots(figsize=(6,4),constrained_layout=True)
    ax.hist(errors,bins=35,color="tab:blue",alpha=.85)
    ax.set(xlabel="Validation relative L2",ylabel="Samples",title=f"All {len(errors)} held-out samples")
    fig.savefig(output/"error_distribution.png",dpi=200)
    plt.close(fig)
    save_json(output/"visualization_manifest.json",dict(figures=["training_curves.png","training_curves.pdf","validation_examples.png","validation_examples.pdf","error_distribution.png"],
              selected_validation_indices=chosen,sample_count=len(errors),completed_epochs=len(rows),metrics=metrics))
    print(f"VISUALIZATION COMPLETED: {output} ({len(rows)} epochs, {len(errors)} validation samples)",flush=True)


if __name__ == "__main__":
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config",default="config/config.yaml")
    args=parser.parse_args()
    render(load_config(args.config))