"""Diagnostic-only rectangular adapter to CPUREF-0.3. No production changes."""
import sys,json,time,argparse,hashlib
from pathlib import Path
import numpy as np
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--repo',type=Path,required=True,help='Candidate checkout at 43c527b')
parser.add_argument('--n',type=int,choices=(32,64),default=32)
parser.add_argument('--output-dir',type=Path,required=True)
args=parser.parse_args()
ROOT=args.repo.resolve()
sys.path[:0]=[str(ROOT/'src'),str(ROOT/'tests/validation/multigrid'),str(ROOT/'tests/validation/taylor_green'),str(ROOT/'tests/unit')]
import cpu_reference as cpu
import oracle
from run_case import provenance
from test_projected_explosion import reconstruct
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.refinement import FixedPatchSet2D
from CelerisLab.lbm.refinement_tables import projected_explosion_tables
from CelerisLab.lbm.blocks import BlockLayout2D
class RectSpec(cpu.Spec):
    @property
    def n_fine(self): return 2*self.n_coarse

def setup(n):
    nu=oracle.U*n/100
    cfg=LBMConfig(nx=n,ny=n,viscosity=nu)
    bounds=oracle.rectangles(n,2)[0]
    p=FixedPatchSet2D.from_config(cfg,bounds)
    spec=RectSpec(n_coarse=n,patch_lo=n//8,patch_hi=7*n//8,halo=n//4,tau_coarse=p.tau_coarse,tau_fine=p.tau_fine,collision='srt')
    q,k=np.divmod(p.coalesce_coarse.astype(int),n*n); y,x=np.divmod(k,n)
    rows=np.column_stack((q,y,x)); links=[]
    for qi,yi,xi in rows:
        sx,sy=xi-cpu.CX[qi],yi-cpu.CY[qi]
        if not p.covered[sy,sx]: links.append(cpu.CornerLink(f'{xi}_{yi}_q{qi}',int(qi),(int(sx),int(sy)),(int(xi),int(yi))))
    assert len(links)==8
    maps=cpu.Maps(~p.covered,p.covered,p.fine_kind==0,p.fine_kind==1,p.fine_kind==2,p.fine_kind,np.arange(2*n)//2,(np.arange(2*n)+.5)/2,rows,tuple(links))
    # Verify the CPU coalescence child locations equal actual production tables.
    children=np.array([[qi*(2*n)**2+(2*yi+dy)*2*n+2*xi+dx for dy in (0,1) for dx in (0,1)] for qi,yi,xi in rows])
    np.testing.assert_array_equal(children,p.coalesce_fine)
    layout=BlockLayout2D.build(n,n,8,8,periodic=(True,True))
    neighbors,offsets=projected_explosion_tables(p,layout)
    inverse=np.full(layout.nblocks*np.prod(layout.local_shape),-1,dtype=int)
    inverse[layout.owner.ravel()]=np.arange(n*n)
    neighbors=inverse[neighbors]
    def explode(gc,cv,gf,fv,m,writes):
        # Only actual transfer footprint, including inner and outer ghost slabs.
        gy,gx=np.divmod(p.ghost_fine.astype(int),2*n)
        if mode[0]=='constant': values=gc.reshape(9,-1)[:,p.ghost_coarse]
        else: values=reconstruct(gc,p,neighbors,offsets)[0]
        assert np.isfinite(values).all()
        gf[:,gy,gx]=values; fv[:,gy,gx]=True; writes[gy,gx]+=1
    mode=['constant']; cpu._explode=explode
    return spec,maps,p,mode,nu

def run(n,which):
    spec,m,p,mode,nu=setup(n); mode[0]=which
    masks=[m.coarse_active,m.fine_active]
    arrays=[oracle.initial(mask.shape,l,n,nu) for l,mask in enumerate(masks)]
    for a,mask in zip(arrays,masks):a[:,~mask]=np.nan
    s=cpu.State(*arrays); rows=[]; initmass=None; momentum0=None; start=time.monotonic()
    for tstar in oracle.TIMES:
        step=round(tstar*n/oracle.U)
        while s.tick<2*step: s,_=cpu.advance_coarse_once(spec,m,s)
        arrays=[s.coarse_f,s.fine_f]
        row=oracle.metrics(arrays,masks,step,n,nu,initmass)
        if initmass is None:initmass=row['mass']
        dot=den=res=0.; parts=[]; momentum=np.zeros(2)
        for l,(a,mask) in enumerate(zip(arrays,masks)):
            v=a[:,mask];rho=v.sum(0);u=cpu.CX@v/rho;w=cpu.CY@v/rho
            x,y=oracle.centers(mask.shape,l);ex,ey,_=oracle.exact(x[mask],y[mask],step,n,nu);vol=.25**l
            dot+=vol*np.sum(u*ex+w*ey);den+=vol*np.sum(ex*ex+ey*ey)
            parts.append((u,w,ex,ey,vol));momentum+=vol*np.array([np.sum(cpu.CX@v),np.sum(cpu.CY@v)])
        amp=dot/den
        res=sum(vol*np.sum((u-amp*ex)**2+(w-amp*ey)**2) for u,w,ex,ey,vol in parts)
        if momentum0 is None:momentum0=momentum.copy()
        row.update(step=step,exact_energy_quadrature_factor=den/(2*n*n*row['exact_energy']),amplitude_ratio=amp,amplitude_bias=amp-1,orthogonal_residual_l2=float(np.sqrt(res/den)),energy_signed_relative=row['energy']/row['exact_energy']-1,mass_signed_relative=row['mass']/initmass-1,momentum=momentum.tolist(),momentum_change=(momentum-momentum0).tolist())
        rows.append(row)
    return dict(n=n,explosion=which,seconds=time.monotonic()-start,rows=rows)
if __name__=='__main__':
    a=args
    a.output_dir.mkdir(parents=True,exist_ok=True)
    result=dict(implementation_revision='43c527b',source=provenance(ROOT),diagnostic_input_sha256={name:hashlib.sha256((ROOT/name).read_bytes()).hexdigest() for name in ('tests/validation/multigrid/cpu_reference.py','tests/unit/test_projected_explosion.py')},precision='FP64 physical populations',phase='arrival_f; macro moments invariant under collision, initialization equilibrium',runs=[])
    out=a.output_dir/f'n{a.n}_ab.json'
    for mode in ('constant','projected_linear'):
        r=run(a.n,mode);result['runs'].append(r);out.write_text(json.dumps(result,indent=2)+'\n');print(json.dumps(r),flush=True)