"""Visualize independent Burgers predictions, errors and the actual training log."""
import csv
import json

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from common import parser, load_config, resolve, save_json, update_run_manifest


def main():
    args = parser(__doc__).parse_args()
    cfg = load_config(args.config)
    folder = resolve(cfg['paths']['results'])
    data = np.load(folder / 'test_predictions.npz')
    pred, truth, times, x = (data[k] for k in ['prediction', 'reference', 'times', 'x'])
    if pred.shape != (200, 11, 33) or not np.isfinite(pred).all():
        raise ValueError('Invalid independent test predictions')
    summary = json.loads((folder / 'test_metrics.json').read_text())
    rmse = np.sqrt(np.mean((pred - truth)**2, axis=(1, 2)))
    plot_dir = folder / 'figures'
    plot_dir.mkdir(exist_ok=True)
    dpi = cfg['visualization']['dpi']
    files = []

    def save(fig, name):
        for suffix in ['png', 'pdf']:
            destination = plot_dir / f'{name}.{suffix}'
            fig.savefig(destination, dpi=dpi, bbox_inches='tight')
            files.append(str(destination.relative_to(folder)))
        plt.close(fig)

    selected = list(dict.fromkeys(cfg['visualization']['sample_indices'] + [int(rmse.argmax())]))
    for i in selected:
        uid = int(data['sample_ids'][i])
        fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), constrained_layout=True)
        vlim = max(float(np.max(np.abs(truth[i]))), float(np.max(np.abs(pred[i]))))
        for ax, field, title in zip(axes[:2], [truth[i], pred[i]], ['Reference', 'RONOM']):
            im = ax.pcolormesh(times, x, field.T, shading='nearest', cmap='RdBu_r',
                               vmin=-vlim, vmax=vlim)
            ax.set(title=title, xlabel='t', ylabel='x')
        fig.colorbar(im, ax=axes[:2], label='u(x,t)')
        im = axes[2].pcolormesh(times, x, np.abs(pred[i] - truth[i]).T,
                               shading='nearest', cmap='magma')
        axes[2].set(title='Absolute error', xlabel='t', ylabel='x')
        fig.colorbar(im, ax=axes[2], label='|prediction - reference|')
        fig.suptitle(f'Independent test case {uid}, trajectory RMSE = {rmse[i]:.3e}')
        save(fig, f'test_case_{uid}_fields')
        fig, axes = plt.subplots(1, 3, figsize=(12, 3.2), constrained_layout=True)
        for ax, j in zip(axes, [0, 5, 10]):
            ax.plot(x, truth[i, j], color='black', label='Reference')
            ax.plot(x, pred[i, j], '--o', ms=3, color='tab:blue', label='RONOM')
            ax.set(title=f't = {times[j]:.1f}', xlabel='x', ylabel='u(x,t)')
            ax.grid(alpha=.2)
        axes[0].legend()
        fig.suptitle(f'Independent test case {uid}')
        save(fig, f'test_case_{uid}_profiles')

    epochs = np.array([], dtype=int)
    if (folder / 'training_history.csv').exists():
        with (folder / 'training_history.csv').open() as f:
            history = list(csv.DictReader(f))
        epochs = np.array([int(r['epoch']) for r in history])
        if len(epochs) != 5000 or not np.array_equal(epochs, np.arange(1, 5001)):
            raise ValueError('Expected one complete formal training history, epochs 1 through 5000')
        val_rows = [r for r in history if r['validation_mse']]
        fig, ax = plt.subplots(figsize=(7, 4))
        ax.semilogy(epochs, [float(r['train_mse']) for r in history], label='Training MSE')
        ax.semilogy([int(r['epoch']) for r in val_rows],
                    [float(r['validation_mse']) for r in val_rows], 'o-', label='Validation MSE')
        ax.axvline(summary['checkpoint_epoch'], color='grey', ls=':', label='Selected epoch')
        ax.set(xlabel='Epoch', ylabel='MSE', title='Formal 5000-epoch training')
        ax.legend(); ax.grid(alpha=.2)
        save(fig, 'training_curve')
    else:
        print('Training history is absent in this results directory; generating independent test figures.', flush=True)
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), constrained_layout=True)
    axes[0].plot(times, summary['rmse_by_time'], 'o-')
    axes[0].set(xlabel='t', ylabel='RMSE across independent test cases and x', title='Error by time')
    axes[1].hist(rmse * 1000, bins=25, edgecolor='white')
    axes[1].set(xlabel='Trajectory RMSE (× 10⁻³)', ylabel='Test cases', title='200 independent trajectories')
    save(fig, 'test_error_summary')
    comparison = folder / 'table1a_comparison.csv'
    with comparison.open('w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['source', 'split', 'samples', 'rmse_mean', 'rmse_std_ddof1', 'units'])
        writer.writerow(['paper_Table1a_RONOM', 'paper_test', 200, .00134, .00162, 'raw'])
        writer.writerow(['this_run', 'independent_test_800_999', 200,
                         summary['trajectory_rmse_mean'], summary['trajectory_rmse_std_ddof1'], 'raw'])
    save_json(folder / 'visualization_manifest.json', {
        'figures': files, 'sample_ids': [int(data['sample_ids'][i]) for i in selected],
        'dpi': dpi, 'test_shape': list(pred.shape), 'training_epochs': len(epochs) if len(epochs) else None,
        'colormaps': {'field': 'RdBu_r', 'absolute_error': 'magma'},
        'note': 'Validation selected weights. Test figures and metrics are post-training only.'})
    update_run_manifest(cfg, status='completed', figure_files=len(files),
                        visualization_manifest='visualization_manifest.json')
    print(f'VISUALIZATION COMPLETE: {len(files)} figures in {plot_dir}', flush=True)


if __name__ == '__main__':
    main()