"""Publication figures generated exclusively from saved numerical results."""
from pathlib import Path
from fractions import Fraction
import json
import os
import tempfile
import numpy as np
os.environ.setdefault('MPLCONFIGDIR',str(Path(tempfile.gettempdir())/'reversible_ca_matplotlib'))
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter


COLORS=['#166496','#d56b2c','#278167','#9c4b8c']


def setup():
    plt.rcParams.update({'font.family':'serif','font.serif':['Computer Modern Roman'],
        'text.usetex':True,
        'text.latex.preamble':r'\usepackage[T1]{fontenc}\usepackage{lmodern}',
        'font.size':11,'axes.labelsize':11,'axes.titlesize':11,'legend.fontsize':10,
        'xtick.labelsize':10,'ytick.labelsize':10,'axes.spines.top':False,
        'axes.spines.right':False,'lines.linewidth':1.5,
        'pdf.fonttype':42,'ps.fonttype':42,'savefig.dpi':180})


def finish(fig,folder,name):
    folder.mkdir(parents=True,exist_ok=True)
    fig.savefig(folder/(name+'.pdf'),bbox_inches='tight')
    fig.savefig(folder/(name+'.png'),bbox_inches='tight',dpi=180)
    plt.close(fig)


def read_json(path):return json.loads(path.read_text())


def convergence_plot(out,folder):
    s=read_json(out/'convergence'/'summary.json');rows=s['rows']
    h=np.array([r['h'] for r in rows])
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    for name,label,color in [('max_l2',r'$L^2$ displacement',COLORS[0]),
                             ('max_h1',r'$H_h^1$ displacement',COLORS[1]),
                             ('max_time_difference_l2',r'$L^2$ time difference',COLORS[2])]:
        axes[0].loglog(h,[r[name] for r in rows],'o-',color=color,label=label,ms=4)
    axes[0].loglog(h,.20*h*h,'--',color='0.45',label=r'$C h^2$')
    axes[0].set(xlabel='Grid spacing $h$',ylabel=r'Maximum error, $0\leq t\leq4$',
                title='(a) Exact nonlinear\ntravelling wave')
    axes[0].legend(frameon=False,loc='upper left')
    axes[0].set_ylim(top=.04)
    axes[0].set_xticks(h,[f'{value:g}' for value in h])
    axes[0].xaxis.set_minor_formatter(NullFormatter())
    axes[0].grid(True,which='both',alpha=.16)
    for color,r in zip(COLORS,rows):
        with np.load(out/'convergence'/f'grid_{r["sites"]}.npz',allow_pickle=False) as z:
            x=np.arange(r['sites'])*r['h']
            error=(z['fields'][-1,0]-z['exact'][-1])/r['h']**2
        axes[1].plot(x,error,color=color,label=f"$J={r['sites']}$")
    axes[1].set(xlabel='$x$',ylabel=r'$(u_h-u_{\rm exact})/h^2$ at $t=4$',
                title='(b) Scaled displacement\nerrors')
    axes[1].legend(frameon=False,ncol=2)
    axes[1].set_ylim(top=.15)
    axes[1].grid(alpha=.16)
    finish(fig,folder,'convergence')


def spectroscopy_plot(out,folder):
    s=read_json(out/'spectroscopy'/'summary.json');rows=s['rows'];speeds=s['speeds']
    finest=max(r['sites'] for r in rows)
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    curve=np.linspace(0,.82,250)
    for branch,mu2,color,marker in [('u',0,COLORS[0],'o'),('v',1,COLORS[1],'s')]:
        selected=[r for r in rows if r['sites']==finest and r['branch']==branch]
        axes[0].plot(curve,np.sqrt(curve**2+mu2),color=color,label=f'PDE {branch} branch')
        axes[0].plot([r['k'] for r in selected],[r['omega_timefit'] for r in selected],
                     marker,ms=4,mfc='white',color=color,label=f'Time fit, {branch}')
        speed=[r for r in speeds if r['branch']==branch]
        axes[1].plot([r['h'] for r in speed],[r['speed_fit']-1 for r in speed],
                     marker+'-',color=color,ms=4,label=f'{branch} branch')
    axes[0].set(xlabel='Physical wavenumber $k$',ylabel=r'Frequency $\omega$',
                title=f'(a) Independent temporal fits\n($J={finest}$)')
    axes[0].legend(frameon=False,loc='upper left')
    axes[0].set_ylim(top=1.85)
    axes[1].axhline(0,color='0.5',ls='--',lw=1)
    axes[1].set(xlabel='Grid spacing $h$',ylabel=r'Fitted speed $c_{\rm fit}-1$',
                title='(b) Approach to\na common speed')
    axes[1].ticklabel_format(axis='y',style='sci',scilimits=(0,0))
    axes[1].legend(frameon=False)
    spacings=sorted(set(r['h'] for r in speeds))
    axes[1].set_xticks(spacings,[f'{value:g}' for value in spacings])
    for ax in axes:ax.grid(alpha=.16)
    finish(fig,folder,'spectroscopy')


def collision_plot(out,folder):
    s=read_json(out/'collision'/'summary.json');principal=s['principal_label']
    with np.load(out/'collision'/(principal+'.npz'),allow_pickle=False) as z:
        times=z['times'];u=z['u'];v=z['v'];vl2=z['v_l2']
    meta=read_json(out/'collision'/(principal+'.json'));length=float(Fraction(meta['model']['length']))
    fig,axes=plt.subplots(2,2,figsize=(6.3,5.7),layout='constrained')
    extent=(0,length,times[0],times[-1])
    for ax,field,title in [(axes[0,0],u,'(a) Incoming and outgoing\n$u$ packets'),
                            (axes[0,1],v,'(b) Generated $v$ response')]:
        scale=np.max(np.abs(field))
        im=ax.imshow(field,extent=extent,origin='lower',aspect='auto',cmap='RdBu_r',
                     vmin=-scale,vmax=scale,interpolation='nearest')
        ax.set(xlabel='$x$',ylabel='$t$',title=title)
        fig.colorbar(im,ax=ax,pad=.02,fraction=.045)
    axes[1,0].plot(times,vl2,color=COLORS[0],label='Both packets')
    for content,color,style in [('left',COLORS[1],'-'),('right',COLORS[2],'--')]:
        with np.load(out/'collision'/(principal.replace('both',content)+'.npz'),allow_pickle=False) as z:
            axes[1,0].plot(z['times'],z['v_l2'],style,color=color,label=f'{content.capitalize()} only')
    axes[1,0].plot(times,np.zeros_like(times),':',color='0.3',label=r'Linear / $\varepsilon=0$')
    axes[1,0].set(xlabel='$t$',ylabel=r'$\|v(t)\|_{L_h^2}$',title='(c) Continuous interaction\nand controls')
    axes[1,0].legend(frameon=False,loc='upper left')
    axes[1,0].set_ylim(top=1.8*np.max(vl2))
    with np.load(out/'collision'/'collision_excess.npz',allow_pickle=False) as z:
        axes[1,1].plot(z['times'],z['u_excess_l2'],color=COLORS[0],label='$u$ excess')
        axes[1,1].plot(z['times'],z['v_excess_l2'],color=COLORS[1],label='$v$ excess')
    axes[1,1].set(xlabel='$t$',ylabel=r'$\|w_{LR}-w_L-w_R\|_{L_h^2}$',
                  title='(d) Collision contribution\nbeyond single packets')
    axes[1,1].legend(frameon=False)
    for ax in axes[1]:ax.grid(alpha=.16)
    finish(fig,folder,'collision')


def stability_plot(out,folder):
    s=read_json(out/'collision'/'summary.json')
    principal=next(r for r in s['runs'] if r['label']==s['principal_label'])
    selected=[r for r in s['runs'] if r['epsilon']==principal['epsilon'] and r['content']=='both' and not r['linear']]
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    for color,r in zip(COLORS,selected):
        with np.load(out/'collision'/(r['label']+'.npz'),allow_pickle=False) as z:
            e=z['diagnostic_energy'];et=z['energy_times']
            axes[0].plot(et,(e-e[0])/e[0],color=color,label=f"$J={r['sites']}$")
            axes[1].semilogy(z['times'],np.maximum(z['uv_total_power_normalized'],1e-30),
                            color=color,label=f"$J={r['sites']}$")
    axes[0].set(xlabel='$t$',ylabel=r'$[E_d(t)-E_d(\tau)]/E_d(\tau)$',
                title='(a) Continuum energy\ndiagnostic')
    axes[0].ticklabel_format(axis='y',style='sci',scilimits=(0,0))
    axes[1].set(xlabel='$t$',ylabel='UV field power /\ninitial total field power',
                title='(b) High-frequency content\n'+r'($|kh|>\pi/2$)')
    for ax in axes:ax.legend(frameon=False);ax.grid(alpha=.16)
    finish(fig,folder,'stability')


def epsilon_plot(out,folder):
    s=read_json(out/'collision'/'summary.json')
    principal=next(r for r in s['runs'] if r['label']==s['principal_label'])
    selected=sorted([r for r in s['runs'] if r['sites']==principal['sites'] and r['epsilon']>0
                     and r['content']=='both' and not r['linear']],key=lambda r:r['epsilon'])
    eps=np.array([r['epsilon'] for r in selected]);norm=np.array([r['max_v_l2'] for r in selected])
    energy=np.array([r['max_v_quadratic_energy'] for r in selected])
    cross=sorted([r for r in s['cross_excess'] if r['sites']==principal['sites']],key=lambda r:r['epsilon'])
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    axes[0].loglog(eps,norm,'o-',color=COLORS[0],label=r'$\max_t\|v\|_2$')
    axes[0].loglog(eps,[r['max_v_excess'] for r in cross],'s-',color=COLORS[1],label=r'Collision excess in $v$')
    axes[0].loglog(eps,norm[-1]*eps/eps[-1],'--',color='0.4',label=r'$C\varepsilon$')
    axes[0].set(xlabel=r'Asymmetry $\varepsilon$',ylabel='Maximum response amplitude',
        title=f"(a) Amplitude slope {s['epsilon_amplitude_slope']:.4f}")
    axes[1].loglog(eps,energy,'o-',color=COLORS[2],label='Quadratic branch energy')
    axes[1].loglog(eps,energy[-1]*(eps/eps[-1])**2,'--',color='0.4',label=r'$C\varepsilon^2$')
    axes[1].set(xlabel=r'Asymmetry $\varepsilon$',ylabel=r'$\max_t E_{v,\mathrm{quad}}$',
        title=f"(b) Quadratic-energy slope {s['epsilon_quadratic_energy_slope']:.4f}")
    for ax in axes:
        ax.legend(frameon=False);ax.grid(True,which='both',alpha=.16)
        ax.set_xticks(eps,[f'{value:g}' for value in eps]);ax.xaxis.set_minor_formatter(NullFormatter())
    axes[0].legend(frameon=False,loc='upper left')
    axes[0].set_ylim(top=3*np.max(norm))
    finish(fig,folder,'epsilon_scaling')


def source_plot(out,folder):
    s=read_json(out/'source_check'/'summary.json')
    cs=read_json(out/'collision'/'summary.json')
    cm=read_json(out/'collision'/(cs['principal_label']+'.json'))['model']
    length=float(Fraction(cm['length']))
    with np.load(out/'source_check'/'first_variation.npz',allow_pickle=False) as z:
        times=z['times'];v1=z['v1'];h=length/v1.shape[1]
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    axes[0].plot(times,np.sqrt(h*np.sum(v1*v1,axis=1)),color='0.2',lw=2.4,label=r'Independent $v_1$ PDE')
    paths=sorted((out/'source_check').glob('epsilon_*.npz'))
    for color,path in zip(COLORS,paths):
        parts=path.stem.split('_')[-1].split('over');eps=float(Fraction(int(parts[0]),int(parts[1])))
        with np.load(path,allow_pickle=False) as z:
            norms=np.sqrt(h*np.sum(z['v']**2,axis=1))/eps
        axes[0].plot(times,norms,'--',color=color,label=rf'$v_\varepsilon/\varepsilon$, $\varepsilon={eps:g}$')
    axes[0].set(xlabel='$t$',ylabel=r'$L_h^2$ norm',title='(a) First-order\ncontinuum response')
    axes[0].legend(frameon=False,loc='upper left')
    axes[0].set_ylim(top=1.8*max(np.max(line.get_ydata()) for line in axes[0].lines))
    rows=sorted(s['rows'],key=lambda r:r['epsilon']);eps=np.array([r['epsilon'] for r in rows])
    error=np.array([r['max_v_over_epsilon_minus_v1_l2'] for r in rows])
    axes[1].loglog(eps,error,'o-',color=COLORS[0],label='Independent PDE comparison')
    axes[1].loglog(eps,error[-1]*(eps/eps[-1])**2,'--',color='0.4',label=r'$C\varepsilon^2$')
    axes[1].set(xlabel=r'$\varepsilon$',ylabel=r'$\max_t\|v_\varepsilon/\varepsilon-v_1\|_2$',
                title=f"(b) Remainder slope {s['epsilon_remainder_order']:.4f}")
    axes[1].legend(frameon=False,loc='upper left')
    axes[1].set_ylim(top=3*np.max(error))
    axes[1].set_xticks(eps,[f'{value:g}' for value in eps]);axes[1].xaxis.set_minor_formatter(NullFormatter())
    for ax in axes:ax.grid(True,which='both',alpha=.16)
    finish(fig,folder,'source_response')


def causal_plot(out,folder):
    with np.load(out/'causal'/'support.npz',allow_pickle=False) as z:
        support=z['support'];times=z['times'];front=z['front_positions']
    fig,axes=plt.subplots(1,2,figsize=(6.3,3.0),layout='constrained')
    axes[0].imshow(support,extent=(-64.5,64.5,-.25,times[-1]+.25),origin='lower',
                   aspect='auto',cmap='Blues',interpolation='nearest',vmin=0,vmax=1)
    for sign in [-1,1]:axes[0].plot(sign*2*times,times,'--',color=COLORS[1],lw=1.2)
    axes[0].set(xlim=(-44,44),xlabel='Distance from initial site',ylabel='Physical time',
                title='(a) Exact register support\nand radius bound')
    axes[1].plot(times,front,color=COLORS[0],label='Observed support radius')
    axes[1].plot(times,2*times,'--',color=COLORS[1],label=r'Locality bound $2t$')
    axes[1].plot(times,times,':',color='0.35',label='PDE characteristic radius $t$')
    axes[1].set(xlabel='Physical time',ylabel='Support radius',title='(b) Microscopic and\ncontinuum scales')
    axes[1].legend(frameon=False,loc='upper left');axes[1].grid(alpha=.16)
    axes[1].set_ylim(top=1.8*np.max(2*times))
    finish(fig,folder,'causal_support')


def make_plots(out,folder):
    out,folder=Path(out),Path(folder);setup()
    choices=[('convergence',convergence_plot),('spectroscopy',spectroscopy_plot),
             ('collision',collision_plot),('collision',stability_plot),('collision',epsilon_plot),
             ('source_check',source_plot),('causal',causal_plot)]
    for name,fn in choices:
        if (out/name/'summary.json').exists():fn(out,folder)


if __name__=='__main__':
    experiment=Path(__file__).resolve().parents[1]
    make_plots(experiment/'results',experiment.parent/'figures')