"""Plot completed training and held-out bubble radius predictions."""
import csv
import json
import sys
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from common import arguments,config,path,write_json


def run(cfg):
    root=path(cfg['paths']['results'])
    a=np.load(root/'predictions.npz',allow_pickle=False)
    with (root/'history.csv').open() as f:history=list(csv.DictReader(f))
    epochs=np.array([int(r['epoch']) for r in history])
    plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False})
    fig,axes=plt.subplots(1,3,figsize=(14,3.8),constrained_layout=True)
    for key,label in [('train_weighted_loss','vRBA weighted loss'),('train_sampled_mse_normalized','Sampled MSE'),('validation_mse_normalized','Validation MSE')]:
        axes[0].semilogy(epochs,[float(r[key]) for r in history],label=label,lw=1)
    axes[0].set(xlabel='Epoch',ylabel='Normalized loss');axes[0].legend()
    axes[1].semilogy(epochs,[float(r['validation_relative_l2_normalized_batch_mean']) for r in history],label='Validation, normalized')
    axes[1].semilogy(epochs,[float(r['validation_relative_l2_physical_global']) for r in history],label='Validation, physical')
    axes[1].set(xlabel='Epoch',ylabel='Relative L2');axes[1].legend()
    axes[2].plot(epochs,[float(r['learning_rate']) for r in history]);axes[2].set(xlabel='Epoch',ylabel='Learning rate')
    fig.savefig(root/'training_curves.png',dpi=180);fig.savefig(root/'training_curves.pdf');plt.close(fig)
    error=a['per_trajectory_relative_l2'];order=np.argsort(error)
    chosen=[int(order[-1]),int(order[len(order)//2]),int(np.random.default_rng(23).integers(len(order)))]
    fig,axes=plt.subplots(3,3,figsize=(14,10),constrained_layout=True)
    for col,idx in enumerate(chosen):
        tt=a['time']*1e6
        for row in [0,1]:
            axes[row,col].plot(tt,a['target'][idx]*1e6,label='Reference',lw=1.8)
            axes[row,col].plot(tt,a['prediction'][idx]*1e6,'--',label='vRBA DeepONet',lw=1.2)
            axes[row,col].set(xlabel='Time (microseconds)',ylabel='Radius (micrometers)')
        axes[0,col].set_title(f'{["Worst","Median","Seeded random"][col]}: trajectory {a["original_indices"][idx]}\nrelative L2 = {error[idx]:.3e}')
        axes[1,col].set_xlim(0,80);axes[1,col].set_title('Initial oscillations')
        axes[0,col].legend()
        axes[2,col].plot(tt,np.abs(a['prediction'][idx]-a['target'][idx])*1e6,color='tab:red')
        axes[2,col].set(xlabel='Time (microseconds)',ylabel='Absolute radius error (micrometers)')
    fig.savefig(root/'bubble_predictions.png',dpi=180);fig.savefig(root/'bubble_predictions.pdf');plt.close(fig)
    fig,axes=plt.subplots(1,2,figsize=(11,4),constrained_layout=True)
    im=axes[0].imshow(np.abs(a['prediction']-a['target'])*1e6,aspect='auto',origin='lower',extent=[0,500,900,1000])
    axes[0].set(xlabel='Time (microseconds)',ylabel='Original test trajectory index',title='Absolute radius error')
    fig.colorbar(im,ax=axes[0],label='Micrometers')
    axes[1].hist(error,bins=20);axes[1].set(xlabel='Per-trajectory relative L2 (physical)',ylabel='Count')
    fig.savefig(root/'test_error_distribution.png',dpi=180);plt.close(fig)
    files=['training_curves.png','training_curves.pdf','bubble_predictions.png','bubble_predictions.pdf','test_error_distribution.png']
    assert all((root/f).stat().st_size>1000 for f in files)
    write_json(root/'visualization_manifest.json',dict(files=files,selected_original_indices=a['original_indices'][chosen].tolist(),
        completed_epochs=int(epochs[-1]),finite_predictions=bool(np.isfinite(a['prediction']).all()),
        test_samples=int(len(error)),plot_units='microseconds, micrometers'))
    print('VISUALIZATION_COMPLETE '+json.dumps(dict(directory=str(root),files=files)),flush=True)


if __name__=='__main__':
    a=arguments(__doc__).parse_args();run(config(a.config))