#!/usr/bin/env python3
"""Scientific field, error and convergence plots from actual run artifacts."""
import argparse
import csv
import json
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.tri as mtri
from scripts.common import load_config, atomic_json


def save(fig, directory, stem):
    for ext in ('png', 'pdf'):
        fig.savefig(directory / f'{stem}.{ext}', dpi=300, bbox_inches='tight')
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--config', default='config/config.yaml')
    args = parser.parse_args()
    cfg, root = load_config(args.config)
    results = root / cfg['paths']['results']
    metrics = json.loads((results / 'metrics.json').read_text())
    manifest = []
    for label in metrics['checkpoints']:
        data = np.load(results / f'predictions_{label}.npz')
        xy, pred, ref = data['xy'], data['prediction_uvp'], data['reference_uvp']
        tri = mtri.Triangulation(xy[:, 0], xy[:, 1])
        fig, axes = plt.subplots(3, 1, figsize=(15, 5.5), constrained_layout=True)
        lower, upper = min(ref[:, 0].min(), pred[:, 0].min()), max(ref[:, 0].max(), pred[:, 0].max())
        values = [ref[:, 0], pred[:, 0], abs(pred[:, 0] - ref[:, 0])]
        titles = ['SU2 reference u', f'Adaptive pseudo-time PINN u ({label}, step {int(data["checkpoint_step"])})', '|u_PINN - u_SU2|']
        for i, (ax, val, title) in enumerate(zip(axes, values, titles)):
            kwargs = dict(cmap='viridis', shading='gouraud')
            if i < 2:
                kwargs.update(vmin=lower, vmax=upper)
            c = ax.tripcolor(tri, val, **kwargs)
            fig.colorbar(c, ax=ax, pad=0.01)
            ax.set(xlim=(0, 15), ylim=(-.5, .5), xlabel='x / H', ylabel='y / H', title=title)
            ax.set_aspect('equal')
        fig.suptitle(f'Re = 800; relative L2(u) = {metrics["checkpoints"][label]["relative_l2_u"]:.6g}')
        save(fig, results, f'velocity_u_{label}')
        manifest += [f'velocity_u_{label}.png', f'velocity_u_{label}.pdf']
        fig, axes = plt.subplots(2, 1, figsize=(15, 4), constrained_layout=True)
        for ax, val, title, cmap in [(axes[0], np.linalg.norm(pred[:, :2], axis=1), 'Predicted speed', 'viridis'),
                                    (axes[1], pred[:, 2], 'Predicted pressure', 'coolwarm')]:
            c = ax.tripcolor(tri, val, cmap=cmap, shading='gouraud')
            fig.colorbar(c, ax=ax, pad=.01)
            ax.set(xlabel='x / H', ylabel='y / H', title=title, xlim=(0, 15), ylim=(-.5, .5))
            ax.set_aspect('equal')
        save(fig, results, f'speed_pressure_{label}')
        manifest += [f'speed_pressure_{label}.png', f'speed_pressure_{label}.pdf']
    rows = list(csv.DictReader((results / 'training_history.csv').open()))
    def series(key):
        valid = [r for r in rows if r.get(key, '')]
        return np.array([float(r['step']) for r in valid]), np.array([float(r[key]) for r in valid])
    fig, axes = plt.subplots(2, 2, figsize=(11, 7), constrained_layout=True)
    x, y = series('loss')
    axes[0, 0].semilogy(x, y, label='Weighted relaxed loss')
    for key in ('raw_u_res', 'raw_v_res', 'raw_p_res'):
        x, y = series(key)
        axes[0, 0].semilogy(x, y, label=key)
    axes[0, 0].legend(fontsize=8)
    x, y = series('relative_l2_u')
    axes[0, 1].semilogy(x, y, label='Actual SU2 relative L2(u)')
    axes[0, 1].axhline(cfg['evaluation']['paper_relative_l2_u'], color='gray', linestyle='--', label='Table 1: 2.32e-3')
    axes[0, 1].legend(fontsize=8)
    for key in ('tau_u', 'tau_v', 'tau_p'):
        x, y = series(key)
        axes[1, 0].semilogy(x, y, label=key)
    axes[1, 0].legend()
    x, y = series('learning_rate')
    axes[1, 1].plot(x, y)
    axes[1, 1].set_ylabel('Learning rate')
    for ax in axes.flat:
        ax.set_xlabel('Actual optimizer iterations')
        ax.grid(alpha=.2)
    save(fig, results, 'training_curves')
    manifest += ['training_curves.png', 'training_curves.pdf']
    summary = json.loads((results / 'training_summary.json').read_text())
    report = ['# BFS Re=800 formal reproduction', '',
              f'Actual completed steps: {summary["completed_steps"]} / {summary["requested_steps"]}.',
              f'Slurm job: {summary["slurm_job_id"]}; host: {summary["host"]}.',
              f'Elapsed training seconds: {summary["elapsed_seconds"]:.1f}.', '',
              '| Selection | Checkpoint step | Relative L2(u) |', '|---|---:|---:|']
    for key, value in metrics['checkpoints'].items():
        report.append(f'| {key} | {value["checkpoint_step"]} | {value["relative_l2_u"]:.8g} |')
    report += ['', 'Paper Table 1 adaptive value: 0.00232. This is a comparison target, not a measured result.',
               'Training used only analytic boundary conditions and resampled PDE collocation points.',
               'The SU2 reference was used only for reporting. Checkpoint selection used an independent unlabeled physics validation set.',
               'Full parameter provenance and paper/repository conflicts are recorded in the reproduction specification.',
               '', 'Figures:', ''] + [f'- {f}' for f in manifest if f.endswith('.png')]
    error_history = results / 'error_history.jsonl'
    if error_history.exists():
        report += ['', 'Execution errors and repairs:', '']
        for line in error_history.read_text().splitlines():
            event = json.loads(line)
            report.append(f'- Job {event["job_id"]}, actual formal steps {event["completed_steps"]}: '
                          f'{event["error"]}. Repair: {event["repair"]}.')
    (results / 'report.md').write_text('\n'.join(report) + '\n')
    atomic_json(results / 'visualization_manifest.json', {'status': 'completed', 'figures': manifest,
                'reference': metrics['reference'], 'dpi': 300, 'interpolation': 'triangulation only for display; metrics use original nodes'})
    print(json.dumps({'event': 'visualization_completed', 'figures': manifest, 'report': str(results / 'report.md')}), flush=True)


if __name__ == '__main__':
    main()