"""Complete, deterministic experiments. All outputs originate in saved trajectories."""
from __future__ import annotations
from fractions import Fraction
from pathlib import Path
import csv
import json
import time
import numpy as np
from scipy.special import ellipj,ellipk
from scipy.optimize import minimize_scalar
from .core import Model,initialize,evolve,as_fraction
from .reference import (spectral_laplacian,reference_solve,unrounded_evolve,
                        branch_fields,field_energy,uv_fraction,discrete_h1,source_reference_solve)
from .tests import run_tests


def write_json(path,value):
    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
    path.write_text(json.dumps(value,ensure_ascii=False,indent=2,allow_nan=False)+'\n')


def write_csv(path,rows):
    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
    with path.open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)


def base_model(cfg,sites,length,**kwargs):
    return Model(sites=sites,length=as_fraction(length),courant=as_fraction(cfg['courant']),
                 delta_factor=cfg['delta_factor'],modulus=int(cfg['modulus']),**kwargs)


def save_run(folder,name,model,data,diagnostics,extras=None):
    folder=Path(folder);folder.mkdir(parents=True,exist_ok=True)
    arrays=dict(data)
    if extras:arrays.update(extras)
    np.savez_compressed(folder/(name+'.npz'),**arrays)
    write_json(folder/(name+'.json'),dict(model=model.to_dict(),diagnostics=diagnostics))


def cn_exact(x,times,length,amplitude=1.,coupling=1.):
    kappa=4*ellipk(.5)/length
    omega=np.sqrt(kappa*kappa+coupling*amplitude*amplitude)
    phase=kappa*np.asarray(x)[None,:]-omega*np.asarray(times)[:,None]
    sn,cn,dn,_=ellipj(phase,.5)
    u=amplitude*cn
    velocity=amplitude*omega*sn*dn
    laplacian=-amplitude*kappa*kappa*cn**3
    return u,velocity,laplacian,kappa,omega


def convergence(cfg,out,tables=None):
    spec=cfg['convergence'];folder=out/'convergence';rows=[]
    previous=None
    for n in spec['sites']:
        model=base_model(cfg,n,spec['length'],lambda_phi=as_fraction(spec['lambda']))
        x=np.arange(n)*float(model.h)
        u0,p0,lap0,kappa,omega=cn_exact(x,[0],float(model.length),spec['amplitude'],spec['lambda'])
        q,p=initialize(model,u0,p0,lap0)
        steps=int(as_fraction(spec['final_time'])/model.tau)
        data,diag=evolve(model,q,p,steps)
        exact,vel,_,_,_=cn_exact(x,data['times'],float(model.length),spec['amplitude'],spec['lambda'])
        err=data['fields'][:,0]-exact
        h,tau=float(model.h),float(model.tau)
        l2=np.sqrt(h*np.sum(err**2,axis=1))
        grad=(np.roll(err,-1,axis=-1)-err)/h
        h1=np.sqrt(h*np.sum(err**2+grad**2,axis=1))
        # Difference of numerical and exact backward time differences.
        dt_error=np.diff(err,axis=0)/tau
        dt_l2=np.sqrt(h*np.sum(dt_error**2,axis=1))
        energy=field_energy(model,data['fields'])
        real=unrounded_evolve(model,u0,p0,lap0,steps)
        quantization_l2=np.sqrt(h*np.sum((data['fields'][:,0]-real[:,0])**2,axis=1))
        row=dict(sites=n,h=h,tau=tau,delta=float(model.delta),max_l2=float(l2.max()),
                 final_l2=float(l2[-1]),max_h1=float(h1.max()),max_time_difference_l2=float(dt_l2.max()),
                 max_quantization_l2=float(quantization_l2.max()),
                 energy_relative_variation=float(np.max(np.abs(energy-energy[0]))/energy[0]),
                 order_l2=None if previous is None else float(np.log2(previous['max_l2']/l2.max())),
                 order_h1=None if previous is None else float(np.log2(previous['max_h1']/h1.max())),
                 max_abs_register=diag['max_abs_register'],nowrap_pass=diag['nowrap_pass'],
                 paired_consistent=diag['paired_consistent'])
        rows.append(row);previous=row
        save_run(folder,f'grid_{n}',model,data,diag,dict(exact=exact,l2_error=l2,h1_error=h1,
                 time_difference_l2_error=dt_l2,unrounded_fields=real,quantization_l2=quantization_l2,
                 diagnostic_energy=energy,energy_times=data['times'][1:-1]))
    write_csv(tables/'convergence.csv',rows)
    summary=dict(rows=rows,kappa=float(kappa),omega=float(omega),boost_velocity=float(kappa/omega),
                 exact_solution='u(t,x)=cn(kappa*x-omega*t | 1/2); lambda=1, amplitude=1',
                 final_time=spec['final_time'],initialization='rounded second-order Taylor sample at -tau')
    write_json(folder/'summary.json',summary);return summary


def mode_frequency(series,tau):
    middle=series[1:-1]
    denominator=2*np.vdot(middle,middle).real
    cosine=np.vdot(middle,series[2:]+series[:-2]).real/denominator
    if not -1<cosine<1:
        raise ValueError(f'estimated cosine outside open oscillatory interval: {cosine}')
    omega=float(np.arccos(cosine)/tau)
    residual=series[2:]-2*cosine*middle+series[:-2]
    relative=float(np.linalg.norm(residual)/(np.linalg.norm(series[2:])+1e-30))
    return omega,relative


def temporal_frequency(times,series,guess):
    """Variable-projection temporal sinusoid fit, independent of the recurrence.

    For every trial frequency, fit complex cosine/sine amplitudes by linear LS.
    A bounded scalar search then minimizes the complete trajectory residual.
    """
    times=np.asarray(times,float);series=np.asarray(series,complex)
    def fit_at(omega):
        basis=np.column_stack([np.cos(omega*times),np.sin(omega*times)])
        coeff=np.linalg.lstsq(basis,series,rcond=None)[0]
        residual=series-basis@coeff
        return float(np.vdot(residual,residual).real)
    result=minimize_scalar(fit_at,bounds=(.9*guess,1.1*guess),method='bounded',
                           options={'xatol':1e-13,'maxiter':150})
    if not result.success:raise RuntimeError('temporal frequency fit failed')
    relative=np.sqrt(result.fun/float(np.vdot(series,series).real))
    return float(result.x),float(relative)


def spectroscopy(cfg,out,tables=None):
    spec=cfg['spectroscopy'];folder=out/'spectroscopy';rows=[];speed_rows=[]
    for n in spec['sites']:
        model=base_model(cfg,n,spec['length'],species=2,eta=as_fraction(spec['eta']))
        x=np.arange(n)*float(model.h)
        for branch in ['u','v']:
            mu2=0. if branch=='u' else 2*float(model.eta)
            modes=[m for m in spec['mode_numbers'] if m!=0 or mu2>0]
            q0=np.zeros(n);p0=np.zeros(n);lap0=np.zeros(n)
            for m in modes:
                k=2*np.pi*m/float(model.length)
                q0+=spec['amplitude_per_mode']*np.cos(k*x)
                p0+=spec['amplitude_per_mode']*np.sqrt(k*k+mu2)*np.sin(k*x)
                lap0-=spec['amplitude_per_mode']*k*k*np.cos(k*x)
            sign=1 if branch=='u' else -1
            initial=np.stack([q0,sign*q0])/np.sqrt(2)
            velocity=np.stack([p0,sign*p0])/np.sqrt(2)
            laplacian=np.stack([lap0,sign*lap0])/np.sqrt(2)
            q,p=initialize(model,initial,velocity,laplacian)
            data,diag=evolve(model,q,p,int(as_fraction(spec['final_time'])/model.tau))
            u,v=branch_fields(data['fields']);selected=u if branch=='u' else v
            transforms=np.fft.fft(selected,axis=-1)/n
            branch_rows=[]
            for m in modes:
                k=2*np.pi*m/float(model.length)
                measured,residual=mode_frequency(transforms[:,m],float(model.tau))
                fitted,fitres=temporal_frequency(data['times'],transforms[:,m],measured)
                middle=len(data['times'])//2
                first,_=temporal_frequency(data['times'][:middle+1],transforms[:middle+1,m],measured)
                second,_=temporal_frequency(data['times'][middle:],transforms[middle:,m],measured)
                ql=4*float(model.courant)**2*np.sin(k*float(model.h)/2)**2+float(model.tau)**2*mu2
                lattice=2*np.arcsin(.5*np.sqrt(ql))/float(model.tau)
                continuum=np.sqrt(k*k+mu2)
                row=dict(sites=n,branch=branch,mode=m,k=k,omega_measured=measured,
                         omega_lattice=float(lattice),omega_continuum=float(continuum),
                         relative_lattice_error=float(abs(measured-lattice)/lattice),
                         relative_continuum_error=float(abs(measured-continuum)/continuum),
                         recurrence_relative_residual=residual,
                         omega_timefit=fitted,timefit_relative_residual=fitres,
                         timefit_lattice_absolute_error=float(abs(fitted-lattice)),
                         timefit_lattice_relative_error=float(abs(fitted-lattice)/lattice),
                         timefit_continuum_relative_error=float(abs(fitted-continuum)/continuum),
                         omega_timefit_first_half=first,omega_timefit_second_half=second,
                         time_window_relative_spread=float(abs(first-second)/fitted))
                rows.append(row);branch_rows.append(row)
            fit=np.linalg.lstsq(np.column_stack([np.ones(len(branch_rows)),np.array([r['k'] for r in branch_rows])**2]),
                                np.array([r['omega_timefit'] for r in branch_rows])**2,rcond=None)[0]
            speed_rows.append(dict(sites=n,h=float(model.h),branch=branch,mass_sq_fit=float(fit[0]),
                speed_fit=float(np.sqrt(fit[1])),max_relative_lattice_error=max(r['timefit_lattice_relative_error'] for r in branch_rows),
                max_relative_continuum_error=max(r['timefit_continuum_relative_error'] for r in branch_rows),
                max_timefit_residual=max(r['timefit_relative_residual'] for r in branch_rows),
                max_time_window_relative_spread=max(r['time_window_relative_spread'] for r in branch_rows)))
            save_run(folder,f'{branch}_grid_{n}',model,data,diag,dict(projected_modes=transforms[:,modes],mode_numbers=np.array(modes,dtype=np.int64)))
    write_csv(tables/'spectroscopy.csv',rows)
    write_csv(tables/'speeds.csv',speed_rows)
    summary=dict(rows=rows,speeds=speed_rows,final_time=spec['final_time'],
                 frequency_estimator='Complete complex trajectory sinusoid least squares with free frequency/amplitudes; recurrence used only for bounded-search initialization and implementation diagnostic',
                 window_check='Independent fits to first and second halves; lowest u mode covers only half a period over the full interval',
                 interpretation='Quantized linear trajectories approach the common-speed continuum branches under refinement; finite observation windows do not establish an exact automaton dispersion relation.')
    write_json(folder/'summary.json',summary);return summary


def periodic_packet(x,center,width,carrier,amplitude,length):
    value=np.zeros_like(x);d1=np.zeros_like(x);d2=np.zeros_like(x)
    for shift in (-length,0,length):
        z=x-center+shift
        envelope=amplitude*np.exp(-.5*(z/width)**2)
        co,si=np.cos(carrier*z),np.sin(carrier*z)
        value+=envelope*co
        d1+=envelope*(-z/width**2*co-carrier*si)
        d2+=envelope*((z*z/width**4-1/width**2-carrier**2)*co+2*carrier*z/width**2*si)
    return value,d1,d2


def collision_initial(model,spec,content='both'):
    x=np.arange(model.sites)*float(model.h)
    k=2*np.pi*spec['carrier_mode']/float(model.length)
    left=periodic_packet(x,spec['left_center'],spec['packet_width'],k,spec['packet_amplitude'],float(model.length))
    right=periodic_packet(x,spec['right_center'],spec['packet_width'],k,spec['packet_amplitude'],float(model.length))
    if content=='left':right=tuple(np.zeros_like(x) for _ in range(3))
    if content=='right':left=tuple(np.zeros_like(x) for _ in range(3))
    u=left[0]+right[0];ut=-left[1]+right[1];lap=left[2]+right[2]
    return np.stack([u,u])/np.sqrt(2),np.stack([ut,ut])/np.sqrt(2),np.stack([lap,lap])/np.sqrt(2)


def collision(cfg,out,tables=None):
    spec=cfg['collision'];folder=out/'collision';folder.mkdir(parents=True,exist_ok=True)
    principal=spec['principal_sites'];base_lambda=as_fraction(spec['lambda']);eps_values=list(map(as_fraction,spec['epsilon']))
    principal_eps=eps_values[0]
    runs={};records=[]
    cases=[]
    for n in spec['sites']:
        for content in ['both','left','right']:
            cases.append((n,principal_eps,content,False))
    for eps in eps_values[1:]:
        for content in (['both','left','right'] if eps else ['both']):
            cases.append((principal,eps,content,False))
    cases.append((principal,-principal_eps,'both',False))
    cases.append((principal,Fraction(0),'both',True))
    for n,eps,content,linear in cases:
        model=base_model(cfg,n,spec['length'],species=2,eta=as_fraction(spec['eta']),
            lambda_phi=Fraction(0) if linear else base_lambda+eps,
            lambda_chi=Fraction(0) if linear else base_lambda-eps,
            g=Fraction(0) if linear else as_fraction(spec['g']))
        initial,velocity,lap=collision_initial(model,spec,content)
        q,p=initialize(model,initial,velocity,lap)
        data,diag=evolve(model,q,p,int(as_fraction(spec['final_time'])/model.tau))
        u,v=branch_fields(data['fields']);h=float(model.h);tau=float(model.tau)
        energy=field_energy(model,data['fields'])
        vnorm=np.sqrt(h*np.sum(v*v,axis=-1))
        uv=uv_fraction(v)
        # Absolute UV power is normalized by the initial total field norm, so a
        # vanishing initial v does not produce a misleading fractional statistic.
        initial_norm_sq=h*np.sum(u[0]**2+v[0]**2)
        k=2*np.pi*np.fft.fftfreq(n)
        vp=np.abs(np.fft.fft(v,axis=-1))**2
        uv_absolute=h/n*vp[:,abs(k)>.5*np.pi+1e-14].sum(axis=-1)/initial_norm_sq
        uv_total=(h/n*np.abs(np.fft.fft(data['fields'],axis=-1))**2).sum(axis=1)
        uv_total=uv_total[:,abs(k)>.5*np.pi+1e-14].sum(axis=-1)/initial_norm_sq
        vt=(v[2:]-v[:-2])/(2*tau);vx=(np.roll(v[1:-1],-1,axis=-1)-v[1:-1])/h
        vquadratic=.5*h*np.sum(vt*vt+vx*vx+2*float(model.eta)*v[1:-1]**2,axis=-1)
        label=f'n{n}_eps{eps.numerator}over{eps.denominator}_{content}'+('_linear' if linear else '')
        record=dict(label=label,sites=n,epsilon=float(eps),content=content,linear=linear,
            max_v_l2=float(vnorm.max()),final_v_l2=float(vnorm[-1]),max_v_quadratic_energy=float(vquadratic.max()),
            final_v_quadratic_energy=float(vquadratic[-1]),
            first_energy_time=float(data['times'][1]),last_energy_time=float(data['times'][-2]),
            max_uv_v_fraction_after_t1=float(uv[data['times']>=1].max()),
            max_uv_v_power_normalized=float(uv_absolute.max()),max_uv_total_power_normalized=float(uv_total.max()),
            diagnostic_energy_relative_variation=float(np.max(abs(energy-energy[0]))/energy[0]),
            final_u_l2=float(np.sqrt(h*np.sum(u[-1]**2))),max_field=float(np.max(abs(data['fields']))),
            max_abs_register=diag['max_abs_register'],nowrap_pass=diag['nowrap_pass'],paired_consistent=diag['paired_consistent'])
        records.append(record)
        save_run(folder,label,model,data,diag,dict(u=u,v=v,v_l2=vnorm,v_quadratic_energy=vquadratic,
            diagnostic_energy=energy,energy_times=data['times'][1:-1],uv_v_fraction=uv,uv_v_power_normalized=uv_absolute,
            uv_total_power_normalized=uv_total))
        runs[(n,eps,content,linear)]=(data,model,u,v,record)
    cross_records=[];excess_runs={}
    for n,eps,content,linear in cases:
        if content!='both' or linear or (n,eps,'left',False) not in runs:continue
        full=runs[(n,eps,'both',False)]
        left=runs[(n,eps,'left',False)];right=runs[(n,eps,'right',False)]
        data,model,u,v,record=full;h=float(model.h)
        du=u-left[2]-right[2];dv=v-left[3]-right[3]
        excess_u=np.sqrt(h*np.sum(du*du,axis=-1));excess_v=np.sqrt(h*np.sum(dv*dv,axis=-1))
        arrays=dict(times=data['times'],u_excess=du,v_excess=dv,u_excess_l2=excess_u,v_excess_l2=excess_v)
        label=f'collision_excess_n{n}_eps{eps.numerator}over{eps.denominator}'
        np.savez_compressed(folder/(label+'.npz'),**arrays)
        if n==principal and eps==principal_eps:
            np.savez_compressed(folder/'collision_excess.npz',**arrays)
        excess_runs[(n,eps)]=arrays
        cross_records.append(dict(sites=n,epsilon=float(eps),max_u_excess=float(excess_u.max()),
            max_v_excess=float(excess_v.max()),final_v_excess=float(excess_v[-1]),
            initial_u_excess=float(excess_u[0]),initial_v_excess=float(excess_v[0])))
    full=runs[(principal,principal_eps,'both',False)]
    principal_excess=excess_runs[(principal,principal_eps)]
    negative=runs[(principal,-principal_eps,'both',False)]
    epsilon_reflection_exact=bool(np.array_equal(full[0]['Q'][:,::-1],negative[0]['Q']) and
                                  np.array_equal(full[0]['P'][:,::-1],negative[0]['P']))
    if not epsilon_reflection_exact:raise AssertionError('epsilon reflection did not interchange species exactly')
    refinement=[]
    ordered=list(spec['sites'])
    for n in ordered[:-1]:
        coarse=runs[(n,principal_eps,'both',False)]
        fine_n=ordered[ordered.index(n)+1];fine=runs[(fine_n,principal_eps,'both',False)]
        ratio=fine_n//n
        difference=coarse[0]['fields']-fine[0]['fields'][::ratio,:,::ratio]
        errors=np.sqrt(float(coarse[1].h)*np.sum(difference*difference,axis=(1,2)))
        verr=coarse[3]-fine[3][::ratio,::ratio]
        verrors=np.sqrt(float(coarse[1].h)*np.sum(verr*verr,axis=1))
        cross_difference=excess_runs[(n,principal_eps)]['v_excess']-excess_runs[(fine_n,principal_eps)]['v_excess'][::ratio,::ratio]
        cross_errors=np.sqrt(float(coarse[1].h)*np.sum(cross_difference**2,axis=1))
        refinement.append(dict(coarse=n,fine=fine_n,max_field_difference=float(errors.max()),
                               final_field_difference=float(errors[-1]),max_v_difference=float(verrors.max()),
                               max_v_cross_excess_difference=float(cross_errors.max())))
    if len(refinement)>=2:
        observed=float(np.log2(refinement[0]['max_field_difference']/refinement[1]['max_field_difference']))
    else:observed=None
    epsilon_records=[r for r in records if r['sites']==principal and r['content']=='both' and not r['linear']]
    epsilon_records=sorted(epsilon_records,key=lambda r:r['epsilon'])
    nonzero=[r for r in epsilon_records if r['epsilon']>0]
    slope_amp=float(np.polyfit(np.log([r['epsilon'] for r in nonzero]),np.log([r['max_v_l2'] for r in nonzero]),1)[0])
    slope_energy=float(np.polyfit(np.log([r['epsilon'] for r in nonzero]),np.log([r['max_v_quadratic_energy'] for r in nonzero]),1)[0])
    epsilon_cross=[r for r in cross_records if r['sites']==principal]
    slope_cross=float(np.polyfit(np.log([r['epsilon'] for r in epsilon_cross]),np.log([r['max_v_excess'] for r in epsilon_cross]),1)[0])
    write_csv(tables/'collision_runs.csv',records)
    write_csv(tables/'collision_refinement.csv',refinement)
    write_csv(tables/'collision_cross_excess.csv',cross_records)
    summary=dict(runs=records,refinement=refinement,refinement_order=observed,
                 epsilon_amplitude_slope=slope_amp,epsilon_quadratic_energy_slope=slope_energy,
                 epsilon_cross_excess_slope=slope_cross,cross_excess=cross_records,
                 cross_excess_refinement_order=(float(np.log2(refinement[0]['max_v_cross_excess_difference']/refinement[1]['max_v_cross_excess_difference'])) if len(refinement)>=2 else None),
                 max_collision_excess_u=float(principal_excess['u_excess_l2'].max()),
                 max_collision_excess_v=float(principal_excess['v_excess_l2'].max()),
                 initial_collision_excess_u=float(principal_excess['u_excess_l2'][0]),
                 initial_collision_excess_v=float(principal_excess['v_excess_l2'][0]),
                 principal_label=full[4]['label'],final_time=spec['final_time'],
                 source_at_v_zero='-epsilon*u^3/2',initial_v_exactly_zero=True,
                 epsilon_reflection_registers_exact=epsilon_reflection_exact,
                 initial_velocity_note='Continuum v_t(0)=0; the staggered increment register includes the required O(tau^2) initial acceleration term',
                 nonlinear_force_active_every_step=True,
                 diagnostic_note='Quadratic branch energy and continuum-type total energy are diagnostics, not exact invariants or particle probabilities.')
    write_json(folder/'summary.json',summary);return summary


def source_check(cfg,out,tables=None):
    """Independent continuum first-variation prediction for branch generation."""
    spec=cfg['collision'];folder=out/'source_check';folder.mkdir(parents=True,exist_ok=True)
    n=spec['principal_sites'];length=float(spec['length']);h=length/n
    base_lambda=as_fraction(spec['lambda']);eta=as_fraction(spec['eta']);g=as_fraction(spec['g'])
    def model_for(nr,eps):
        return base_model(cfg,nr,spec['length'],species=2,eta=eta,
                          lambda_phi=base_lambda+eps,lambda_chi=base_lambda-eps,g=g)
    model=model_for(n,Fraction(0));w0,p0,_=collision_initial(model,spec)
    u0=(w0[0]+w0[1])/np.sqrt(2);ut0=(p0[0]+p0[1])/np.sqrt(2)
    times=np.arange(int(as_fraction(spec['final_time'])/model.tau)+1)*float(model.tau)
    source,meta=source_reference_solve(u0,ut0,length,float(eta),float(base_lambda),float(g),times)
    np.savez_compressed(folder/'first_variation.npz',times=times,u_background=source[:,0],v1=source[:,1],
                        u_background_velocity=source[:,2],v1_velocity=source[:,3])
    write_json(folder/'first_variation.json',meta)
    source_norm=np.sqrt(h*np.sum(source[:,1]**2,axis=1))
    rows=[];principal_reference=None
    for eps in map(as_fraction,spec['epsilon']):
        if eps<=0:continue
        emodel=model_for(n,eps)
        fields,velocity,metadata=reference_solve(emodel,w0,p0,times)
        u,v=branch_fields(fields)
        difference=v/float(eps)-source[:,1]
        error=np.sqrt(h*np.sum(difference**2,axis=1))
        row=dict(epsilon=float(eps),max_v_over_epsilon_minus_v1_l2=float(error.max()),
                 max_relative_to_v1_norm=float(error.max()/source_norm.max()),
                 max_continuum_v_l2=float(np.sqrt(h*np.sum(v*v,axis=1)).max()))
        path=out/'collision'/f'n{n}_eps{eps.numerator}over{eps.denominator}_both.npz'
        if path.exists():
            with np.load(path,allow_pickle=False) as integer_run:
                integer_difference=integer_run['v']-v
                row['max_automaton_v_minus_continuum_v_l2']=float(np.sqrt(h*np.sum(integer_difference**2,axis=1)).max())
        rows.append(row)
        label=f'epsilon_{eps.numerator}over{eps.denominator}'
        np.savez_compressed(folder/(label+'.npz'),times=times,fields=fields,velocities=velocity,u=u,v=v,
                            first_variation_difference_l2=error)
        write_json(folder/(label+'.json'),dict(model=emodel.to_dict(),reference=metadata))
        if principal_reference is None:principal_reference=fields
    coarse_model=model_for(n//2,as_fraction(spec['epsilon'][0]))
    coarse_w,coarse_p,_=collision_initial(coarse_model,spec)
    coarse_ref,coarse_vel,coarse_meta=reference_solve(coarse_model,coarse_w,coarse_p,times)
    spatial_check=np.sqrt(2*h*np.sum((coarse_ref-principal_reference[:,:,::2])**2,axis=(1,2)))
    tight_source,tight_meta=source_reference_solve(u0,ut0,length,float(eta),float(base_lambda),float(g),times,
                                                 rtol=3e-13,atol=3e-14,max_step_factor=.25)
    temporal_check=np.sqrt(h*np.sum((tight_source[:,:2]-source[:,:2])**2,axis=(1,2)))
    np.savez_compressed(folder/'reference_checks.npz',times=times,coarse_reference=coarse_ref,
        coarse_velocities=coarse_vel,tighter_source=tight_source,
        spatial_difference_l2=spatial_check,source_temporal_difference_l2=temporal_check)
    order=float(np.polyfit(np.log([r['epsilon'] for r in rows]),
                          np.log([r['max_v_over_epsilon_minus_v1_l2'] for r in rows]),1)[0])
    summary=dict(rows=rows,epsilon_remainder_order=order,max_v1_l2=float(source_norm.max()),
        continuum_reference_spatial_check_max_l2=float(spatial_check.max()),
        source_reference_temporal_step_halving_check_max_l2=float(temporal_check.max()),
        reference_check_metadata=[coarse_meta,tight_meta],
        equation='v1_tt-v1_xx+[2 eta+(3 lambda-g)u0^2/2] v1=-u0^3/2; v1(0)=v1_t(0)=0',
        interpretation='At fixed finite time v_epsilon/epsilon approaches the independently solved first variation with an O(epsilon^2) remainder.')
    write_csv(tables/'source_check.csv',rows)
    write_json(folder/'summary.json',summary);return summary


def reference_check(cfg,out,tables=None):
    spec=cfg['reference_check'];folder=out/'reference_check';folder.mkdir(parents=True,exist_ok=True)
    def model_for(n):return base_model(cfg,n,spec['length'],species=2,mass_phi_sq=Fraction(1,4),
         mass_chi_sq=Fraction(1,2),eta=Fraction(1,4),lambda_phi=Fraction(3,10),lambda_chi=Fraction(1,10),g=Fraction(1,10))
    def initial(n):
        x=np.arange(n)*spec['length']/n;k=2*np.pi/spec['length']
        u=np.stack([.4*np.cos(k*x)+.1*np.cos(2*k*x),.3*np.sin(k*x)])
        p=np.stack([.15*np.sin(k*x),-.1*np.cos(k*x)])
        return u,p,spectral_laplacian(u,spec['length'])
    n=spec['sites'];model=model_for(n);u,p,lap=initial(n)
    q,pr=initialize(model,u,p,lap);data,diag=evolve(model,q,pr,int(as_fraction(spec['final_time'])/model.tau))
    references=[];metadata=[]
    for nr in spec['reference_sites']:
        mr=model_for(nr);ur,pr,_=initial(nr)
        ref,refvel,meta=reference_solve(mr,ur,pr,data['times'])
        references.append(ref[:,:,::(nr//n)]);metadata.append(meta)
        np.savez_compressed(folder/f'reference_grid_{nr}.npz',times=data['times'],fields=ref,velocities=refvel)
    refined=references[-1];error=data['fields']-refined
    l2=np.sqrt(float(model.h)*np.sum(error**2,axis=(1,2)))
    refdiff=references[0]-references[1]
    check=np.sqrt(float(model.h)*np.sum(refdiff**2,axis=(1,2)))
    # A tighter temporal tolerance on the fine reference checks the ODE solve.
    nr=spec['reference_sites'][-1];mr=model_for(nr);ur,pr,_=initial(nr)
    tight,tightvel,tightmeta=reference_solve(mr,ur,pr,data['times'],rtol=3e-13,atol=3e-14,max_step_factor=.25)
    np.savez_compressed(folder/'reference_tighter_step.npz',times=data['times'],fields=tight,velocities=tightvel)
    temporal=np.sqrt(float(model.h)*np.sum((tight[:,:,::(nr//n)]-refined)**2,axis=(1,2)))
    summary=dict(sites=n,max_l2_error=float(l2.max()),final_l2_error=float(l2[-1]),
                 reference_spatial_check_max_l2=float(check.max()),reference_tolerance_check_max_l2=float(temporal.max()),
                 reference_metadata=metadata+[tightmeta],
                 temporal_check='Tighter tolerance and maximum time step halved; difference is at floating-point roundoff scale')
    save_run(folder,'two_species',model,data,diag,dict(reference_fields=refined,l2_error=l2,
        reference_spatial_difference_l2=check,reference_temporal_difference_l2=temporal))
    write_json(folder/'summary.json',summary);return summary


def causal_and_inverse(cfg,out,tables=None):
    folder=out/'causal';folder.mkdir(parents=True,exist_ok=True)
    model=Model(129,Fraction(129),modulus=65537)
    q=np.zeros((1,129),dtype=object);p=q.copy();q[0,64]=512
    data,diag=evolve(model,q,p,40)
    distance=np.minimum(abs(np.arange(129)-64),129-abs(np.arange(129)-64))
    support=(data['Q'][:,0]!=0)|(data['P'][:,0]!=0)
    front=np.array([distance[row].max() if row.any() else 0 for row in support])
    outside=np.array([np.any(support[t,distance>t]) for t in range(41)])
    if outside.any():raise AssertionError('classical cone violation')
    qr,pr=data['Q'][-1].astype(object),data['P'][-1].astype(object)
    for _ in range(40):qr,pr=model.inverse(qr,pr,True)
    if not(np.array_equal(qr,q) and np.array_equal(pr,p)):raise AssertionError('roundtrip failed')
    save_run(folder,'support',model,data,diag,dict(support=support,front_positions=front,distances=distance))
    summary=dict(outside_cone_count=0,roundtrip_exact=True,steps=40,front_positions=front.tolist(),
                 support_statement='support radius <= n after n steps; saturation is not required',
                 microscopic_front_bound_physical=float(model.h/model.tau))
    write_json(folder/'summary.json',summary);return summary


def run_all(cfg,out,selection='all',tables=None):
    out=Path(out);out.mkdir(parents=True,exist_ok=True)
    tables=Path(tables) if tables is not None else out.parent/'tables'
    functions=dict(convergence=convergence,spectroscopy=spectroscopy,collision=collision,source_check=source_check,
                   reference_check=reference_check,causal=causal_and_inverse)
    summary={};start=time.perf_counter()
    if selection in ('all','tests'):
        summary['tests']=run_tests();write_json(out/'tests.json',summary['tests'])
    for name,fn in functions.items():
        if selection in ('all',name):
            before=time.perf_counter();print(f'Running {name} ...',flush=True)
            summary[name]=fn(cfg,out,tables);summary[name]['runtime_seconds']=time.perf_counter()-before
            print(f'Finished {name} in {summary[name]["runtime_seconds"]:.2f} s',flush=True)
    summary['runtime_seconds']=time.perf_counter()-start
    if selection=='all':write_json(out/'summary.json',summary)
    else:write_json(out/f'run_{selection}.json',summary)
    return summary