"""Render Darcy coefficient/GT/prediction/signed-error panels and diagnostics."""
from __future__ import annotations

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 atomic_json, configuration, read_config


def main():
    args=configuration('Visualize the full HNO Darcy experiment').parse_args()
    c=read_config(args.config)
    output=Path(c['output']['results_dir'])
    data=np.load(output/'predictions.npz')
    metrics=json.loads((output/'metrics.json').read_text())
    with (output/'training_history.csv').open() as f:
        history=list(csv.DictReader(f))
    epochs=np.asarray([int(r['epoch']) for r in history])
    if len(history)!=c['training']['epochs'] or list(epochs)!=list(range(1,c['training']['epochs']+1)):
        raise ValueError('Visualization requires the complete formal training history')
    if len(data['prediction'])!=c['data']['n_eval']:
        raise ValueError('Incomplete inference data')
    dpi=c['visualization']['dpi']
    fig,axes=plt.subplots(1,3,figsize=(14,3.8),constrained_layout=True)
    for name,label in [('train_relative_l2','Training relative L2'),('eval_relative_l2','Evaluation relative L2')]:
        axes[0].semilogy(epochs,[float(r[name]) for r in history],label=label,lw=1)
    axes[0].axvline(metrics['best_epoch'],color='gray',ls=':',label='Best checkpoint')
    axes[0].axhline(metrics['paper_relative_l2'],color='black',ls='--',label='Paper: 0.00446')
    axes[0].set(xlabel='Epoch',ylabel='Mean relative L2');axes[0].legend(fontsize=8)
    axes[1].semilogy(epochs,[float(r['train_loss']) for r in history],label='Total loss / sample')
    axes[1].semilogy(epochs,[float(r['train_derivative']) for r in history],label='Derivative relative loss')
    axes[1].set(xlabel='Epoch',ylabel='Loss');axes[1].legend(fontsize=8)
    axes[2].plot(epochs,[float(r['learning_rate']) for r in history]);axes[2].set(xlabel='Epoch',ylabel='Learning rate')
    fig.savefig(output/'training_curves.png',dpi=dpi);fig.savefig(output/'training_curves.pdf');plt.close(fig)
    chosen=c['visualization']['sample_indices']
    fig,axes=plt.subplots(len(chosen),4,figsize=(13,3*len(chosen)),squeeze=False,constrained_layout=True)
    titles=['Coefficient a','Ground truth u','HNO prediction','Signed error (prediction - GT)']
    for row,index in enumerate(chosen):
        a,u,p=(data[k][index] for k in ['coefficient','target','prediction'])
        residual=p-u
        vmin=min(u.min(),p.min());vmax=max(u.max(),p.max())
        error_bound=max(float(np.abs(residual).max()),1e-12)
        specifications=[(a,'viridis',None,None),(u,'viridis',vmin,vmax),(p,'viridis',vmin,vmax),
                        (residual,'RdBu_r',-error_bound,error_bound)]
        for col,(field,cmap,lo,hi) in enumerate(specifications):
            ax=axes[row,col]
            im=ax.imshow(field,origin='lower',extent=(0,1,0,1),cmap=cmap,vmin=lo,vmax=hi)
            ax.set(xlabel='x',ylabel='y',aspect='equal')
            if row==0:ax.set_title(titles[col],fontsize=10)
            fig.colorbar(im,ax=ax,fraction=.046,pad=.04)
        axes[row,0].set_ylabel(f'Sample {index}\ny')
        axes[row,2].text(.02,.98,f'Rel. L2={data["relative_l2"][index]:.5f}',transform=axes[row,2].transAxes,
                         va='top',fontsize=8,bbox=dict(facecolor='white',alpha=.8,edgecolor='none'))
    fig.savefig(output/'darcy_predictions.png',dpi=dpi);fig.savefig(output/'darcy_predictions.pdf');plt.close(fig)
    fig,ax=plt.subplots(figsize=(6,4),constrained_layout=True)
    ax.hist(data['relative_l2'],bins=24,color='#35668b',edgecolor='white')
    ax.axvline(metrics['mean_relative_l2'],label=f'Mean: {metrics["mean_relative_l2"]:.6f}',color='black')
    ax.set(xlabel='Per-sample relative L2',ylabel='Evaluation samples',title='Darcy: 200 best-checkpoint predictions')
    ax.legend();fig.savefig(output/'error_distribution.png',dpi=dpi);plt.close(fig)
    # Appendix J.3: aggregate Fourier energy over the full 200-case evaluation set.
    residual=data['prediction']-data['target']
    error_energy=np.abs(np.fft.fft2(residual,axes=(-2,-1)))**2
    target_energy=np.abs(np.fft.fft2(data['target'],axes=(-2,-1)))**2
    size=residual.shape[-1];freq=np.fft.fftfreq(size)*size
    yy,xx=np.meshgrid(freq,freq,indexing='ij');bins=np.floor(np.sqrt(xx**2+yy**2)).astype(int)
    spectrum=[]
    for b in range(bins.max()+1):
        select=bins==b;den=target_energy[:,select].sum()
        spectrum.append(float(np.sqrt(error_energy[:,select].sum()/max(den,1e-30))))
    fig,ax=plt.subplots(figsize=(6,4),constrained_layout=True)
    ax.semilogy(np.arange(len(spectrum)),spectrum);ax.set(xlabel='Radial frequency |k|',ylabel='Relative error spectrum')
    fig.savefig(output/'error_spectrum.png',dpi=dpi);plt.close(fig)
    atomic_json(output/'visualization_manifest.json',{'status':'completed','n_eval':len(data['prediction']),
                'epochs':len(history),'sample_indices':chosen,'dpi':dpi,'colormaps':['viridis','RdBu_r'],
                'figures':['training_curves.png','training_curves.pdf','darcy_predictions.png','darcy_predictions.pdf',
                           'error_distribution.png','error_spectrum.png'],'radial_error_spectrum':spectrum})
    print(f'VISUALIZATION_COMPLETED results={output} epochs={len(history)} samples={len(data["prediction"])}',flush=True)


if __name__=='__main__':
    main()