"""Bounded exact-source CUDA plumbing gate; never a flow-accuracy certificate.

Run this driver externally against exact dev with --mode baseline first, then
against the experimental source with --mode candidate --baseline <receipt>.
Only scalar receipts/hashes and actual compiler sidecars are retained. Whole-step
mass/momentum are diagnostics: no verified FP32 conservation bound is claimed.
"""
import argparse
from contextlib import contextmanager
from dataclasses import replace
import hashlib
import json
import os
import shutil
from pathlib import Path
import subprocess
import sys

import numpy as np

BASE = 'ed9d8c937d446b854d6963bd92d718bb4fb98b96'
PRODUCTION = 'a135252ae78b47b57d7e77d19001700b3bedac8b'
STOPS = (0, 1, 2, 8)
OPERATOR_ATOL, OPERATOR_RTOL = 6e-8, 2e-7
REST_ATOL, REST_RTOL = 3e-7, 1e-6
REPLAY_ATOL = 3e-6
E = np.array([(0,0),(1,0),(-1,0),(0,1),(0,-1),(1,1),(-1,-1),(1,-1),(-1,1)])
W = np.array([4/9]+[1/9]*4+[1/36]*4)


def digest(array):
    """All finite bits (including signed zero) and exact NaN mask; no NaN payloads."""
    array = np.asarray(array)
    if np.isinf(array).any():
        raise AssertionError('infinite population')
    finite = np.isfinite(array)
    h = hashlib.sha256()
    h.update(str((array.shape, array.dtype.str)).encode())
    h.update(np.ascontiguousarray(finite).tobytes())
    h.update(np.ascontiguousarray(array[finite]).tobytes())
    return h.hexdigest()


def independent_explosion(coarse, pair, layout, kind):
    """Coordinate oracle; never reads the production neighbor/offset tables."""
    ny, nx = pair.coarse_shape
    legal = np.zeros((ny,nx), bool)
    legal.ravel()[pair.ghost_coarse] = True
    legal &= layout.owned & ~pair.covered & (kind == 0)
    if not legal.ravel()[pair.ghost_coarse].all():
        raise AssertionError('illegal center donor')
    out = np.empty((9,len(pair.ghost_fine)))
    flat = np.asarray(coarse, np.float64).reshape(9,-1)
    for row, (target, donor) in enumerate(zip(pair.ghost_fine,pair.ghost_coarse)):
        y,x = divmod(int(donor),nx)
        gradient = np.zeros((9,2))
        for axis, (dx,dy) in enumerate(((1,0),(0,1))):
            ends = [(x-dx,y-dy),(x+dx,y+dy)]
            if not layout.periodic[axis] and any(not (0<=xx<nx and 0<=yy<ny) for xx,yy in ends):
                continue
            ends = [(xx%nx,yy%ny) for xx,yy in ends]
            if all(legal[yy,xx] for xx,yy in ends):
                (xm,ym),(xp,yp) = ends
                gradient[:,axis] = (flat[:,yp*nx+xp]-flat[:,ym*nx+xm])/2
        for q in range(1,9):
            gradient[q] -= E[q]*np.dot(E[q],gradient[q])/np.dot(E[q],E[q])
        fy,fx = divmod(int(target),pair.fine_shape[1])
        offset = np.array([pair.fine_origin[0]+(fx+.5)/2-(x+.5),
                           pair.fine_origin[1]+(fy+.5)/2-(y+.5)])
        out[:,row] = flat[:,int(donor)] + gradient@offset
    return out


def default_options(selector):
    # Exact dev has no explosion keyword. Baseline and omitted paths pass none.
    return {} if selector == 'omitted' else {'explosion':selector}


def snapshot(run, raw):
    from CelerisLab.lbm.refinement import ACTIVE, OUTER_GHOST
    stored, public = raw(run), run.download()
    hashes, totals = [], np.zeros(3)
    for level,(values,visible,kind) in enumerate(zip(stored,public,run.hierarchy.kinds)):
        active = kind == ACTIVE
        if not np.isfinite(visible[:,active]).all() or not np.isnan(visible[:,~active]).all():
            raise AssertionError('public ACTIVE/invalid mask mismatch')
        if not np.isnan(values[:,kind==OUTER_GHOST]).all():
            raise AssertionError('OUTER owner core is not poisoned')
        population = visible[:,active].sum(axis=1,dtype=np.float64)*(.25**level)
        totals += np.r_[population.sum(), E.T@population]
        hashes.append(dict(raw=digest(values), public=digest(visible), active=int(active.sum())))
    if not np.isfinite(totals).all():
        raise AssertionError('nonfinite mass/momentum diagnostic')
    return dict(step=run.step_count, levels=hashes, mass_momentum=totals.tolist())


def initialize(run, rest=False):
    arrays, origin, spacing = [], np.zeros(2), 1.
    for index,level in enumerate(run._levels):
        if index:
            origin += spacing*np.asarray(run.hierarchy.pairs[index-1].fine_origin)
            spacing /= 2
        y,x = np.indices(level.shape[1:],dtype=np.float64)
        x,y = origin[0]+spacing*(x+.5),origin[1]+spacing*(y+.5)
        q = np.arange(9)[:,None,None]
        values = np.broadcast_to(W[:,None,None],level.shape).copy()
        if not rest:
            # q-distinct smooth non-equilibrium G; same physical frame on all levels.
            values += 1e-4*np.sin(.17*x+.23*y+q)
        arrays.append(values)
    run.initialize(*arrays)


def operator_case(run, index):
    import pycuda.driver as cuda
    parent,child = run._levels[index:index+2]
    geometry,pair = run.hierarchy.pairs[index],run._pairs[index]
    c = np.random.default_rng(19+index).normal(0,.001,(9,*geometry.coarse_shape))
    if not run._ddf_shifting:
        c += W[:,None,None]
    c = c.astype(np.float32)
    legal = np.zeros(geometry.coarse_shape,bool)
    legal.ravel()[geometry.ghost_coarse] = True
    c[:,~legal] = np.nan
    packed = parent.layout.to_blocked(c)
    output = np.full(child.layout.blocked_shape,np.nan,np.float32)
    cuda.memcpy_htod(parent.current,packed)
    cuda.memcpy_htod(child.current,output)
    calls, hop = [],run._hop
    try:
        run._hop = lambda *a,**kw: calls.append(a)
        run._advance(index+1,parent.current,block=run._block,stream=run._stream)
        run._stream.synchronize()
    finally:
        run._hop = hop
    if calls != [(index+1,child.current,child.spare),(index+1,child.spare,child.current)]:
        raise AssertionError('explosion changed two-hop launch contract')
    cuda.memcpy_dtoh(output,child.current)
    target = child.layout.owner.ravel()[geometry.ghost_fine]
    slots = np.array([child.layout.slot(target,q) for q in range(9)])
    actual = output.ravel()[slots]
    expected = independent_explosion(c,geometry,parent.layout,run.hierarchy.kinds[index])
    constant = c.reshape(9,-1)[:,geometry.ghost_coarse]
    if not np.isfinite(actual).all() or np.max(abs(expected-constant)) <= 1e-5:
        raise AssertionError('operator must exercise finite nonzero projected slopes')
    np.testing.assert_allclose(actual,expected,atol=OPERATOR_ATOL,rtol=OPERATOR_RTOL)
    untouched = np.ones(output.size,bool); untouched[slots.ravel()] = False
    if not np.isnan(output.ravel()[untouched]).all():
        raise AssertionError('explosion wrote outside its targets')
    return dict(pair=index,rows=len(target),parent_stride=pair.parent_cells,
                child_stride=int(np.prod(child.layout.local_shape)),
                max_error_by_q=np.max(abs(actual-expected),axis=1).tolist(),
                nonzero_projection=True,poisoned_nondonors=True,untouched_preserved=True)


def replay_two_hops(values, pair, layout, kind):
    """Existing FP64 two-hop algebra, using the independent coordinate oracle."""
    from tests.unit.test_projected_explosion import collide
    coarse,fine = (a.copy() for a in values)
    fine.reshape(9,-1)[:,pair.ghost_fine] = independent_explosion(coarse,pair,layout,kind)
    for _ in range(2):
        fine = np.array([np.roll(fine[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)])
        fine[:,pair.fine_kind==2] = np.nan
        active = pair.fine_kind==0
        fine[:,active] = collide(fine[:,active],pair.tau_fine)
    out = np.array([np.roll(coarse[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)])
    out.reshape(-1)[pair.coalesce_coarse] = fine.reshape(-1)[pair.coalesce_fine].mean(axis=1)
    out[:,~pair.covered] = collide(out[:,~pair.covered],pair.tau_coarse)
    out[:,pair.covered] = np.nan
    return out,fine


def run_case(ctx, case, shifted, selector, raw, projected=False, replay=False):
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    cfg,refs,blocks = case
    cfg = replace(cfg,velocity=0,ddf_shifting=shifted)
    with FixedPatchStepper(cfg,context=ctx,refinements=refs,level_blocks=blocks,
                           **default_options(selector)) as run:
        operator = [operator_case(run,i) for i in range(len(run._pairs))] if projected else []
        counts = [0]*len(run._levels)
        for index,level in enumerate(run._levels):
            def counted(*a, index=index, kernel=level.kernel, **kw):
                counts[index] += 1
                return kernel(*a,**kw)
            level.kernel = counted
        def check_counts(stop):
            if counts != [stop*2**index for index in range(len(counts))]:
                raise AssertionError('actual level hop counts differ from acoustic schedule')
        rest = []
        if projected:
            initialize(run,rest=True)
            for stop in STOPS:
                run.step(stop-run.step_count)
                check_counts(stop)
                snapshot(run,raw)  # actual public and raw invalid masks at every stop
                maximum = 0.
                for values,kind in zip(run.download(),run.hierarchy.kinds):
                    active = values[:,kind==0]
                    np.testing.assert_allclose(active,np.broadcast_to(W[:,None],active.shape),
                                               atol=REST_ATOL,rtol=REST_RTOL)
                    maximum = max(maximum,float(np.max(abs(active-W[:,None]))))
                rest.append(dict(step=stop,max_absolute_error=maximum))
        counts[:] = [0]*len(counts)
        initialize(run)
        reference = tuple(a.astype(np.float64) for a in run.download()) if replay else None
        snapshots,replay_errors,previous = [],[],0
        for stop in STOPS:
            run.step(stop-run.step_count)
            if replay:
                for _ in range(stop-previous):
                    reference = replay_two_hops(reference,run.hierarchy.pairs[0],
                                               run._levels[0].layout,run.hierarchy.kinds[0])
                errors = []
                for actual,expected,kind in zip(run.download(),reference,run.hierarchy.kinds):
                    np.testing.assert_allclose(actual[:,kind==0],expected[:,kind==0],
                                               atol=REPLAY_ATOL,rtol=0)
                    errors.append(np.max(abs(actual[:,kind==0]-expected[:,kind==0]),axis=1).tolist())
                replay_errors.append(dict(step=stop,max_error_by_level_q=errors))
            check_counts(stop)
            snapshots.append(dict(snapshot(run,raw),observed_level_hops=counts.copy())); previous = stop
        change = (np.array(snapshots[-1]['mass_momentum'])-
                  np.array(snapshots[0]['mass_momentum'])).tolist()
        return dict(operator=operator,snapshots=snapshots,rest_fixed_point=rest,
                    replay_max_errors=replay_errors,mass_momentum_change_diagnostic=change,
                    conservation_status='diagnostic_only_no_FP32_bound')


def compare_snapshots(actual,expected):
    if len(actual) != len(expected):
        raise AssertionError('snapshot count differs')
    for got,want in zip(actual,expected):
        if (got['step'] != want['step'] or got['levels'] != want['levels']
                or got.get('observed_level_hops') != want.get('observed_level_hops')):
            raise AssertionError('owner bits, public values or masks differ')


def cases():
    from CelerisLab.config import LBMConfig
    from tests.unit.test_hierarchy import CFG, FOUR
    scalar = (LBMConfig(nx=16,ny=16,viscosity=.02),
              [(4,4,12,12),(8,8,16,16),(8,8,16,16)],
              [(4,4),(4,8),(6,4),(4,6)])
    four = (CFG,FOUR,[(6,8),(8,6),(12,8),(16,12)])
    concave = (LBMConfig(nx=30,ny=30,viscosity=.1),
               [[(4,4,5,7),(4,5,7,7)]],[(6,5),(4,6)])
    return dict(scalar=scalar,four=four,concave=concave)


def baseline_receipt(receipt, driver_hash, keys):
    if (receipt.get('completed') is not True or receipt.get('dry_run') is not False
            or receipt.get('mode') != 'baseline' or receipt.get('production_revision') != BASE
            or receipt.get('source',{}).get('revision') != BASE
            or receipt.get('source',{}).get('dirty') is not False
            or receipt.get('driver_sha256') != driver_hash
            or receipt.get('stops') != list(STOPS)
            or receipt.get('operator_tolerance') != [OPERATOR_ATOL,OPERATOR_RTOL]
            or receipt.get('rest_tolerance') != [REST_ATOL,REST_RTOL]
            or receipt.get('replay_tolerance') != [REPLAY_ATOL,0]
            or set(receipt.get('default_snapshots',{})) != set(keys)
            or not receipt.get('actual_build_sidecars')):
        raise ValueError('baseline receipt does not match the frozen exact-dev gate')
    return receipt


def remap_case(ctx, shifted, raw):
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    cfg,refs,blocks = cases()['four']
    with FixedPatchStepper(replace(cfg,velocity=0,ddf_shifting=shifted),context=ctx,
            refinements=refs[:2],level_blocks=blocks[:3],explosion='projected_linear') as run:
        initialize(run); run.step(1)
        before = snapshot(run,raw)
        with run.remap_child([(22,22,42,42),(72,68,76,76)]) as clone:
            if clone._explosion != 'projected_linear' or any(p.neighbors is None for p in clone._pairs):
                raise AssertionError('remap lost projected selection')
            if clone.step_count != 1: raise AssertionError('remap lost synchronized count')
            initial = snapshot(clone,raw)
            clone.step(1)
            after = snapshot(clone,raw)
            if clone.step_count != 2: raise AssertionError('remap did not continue count')
        compare_snapshots([snapshot(run,raw)],[before])
        return dict(shifted=shifted,source_unchanged=True,initial=initial,continued=after)


@contextmanager
def capture_builds(report, save, baseline=None):
    from CelerisLab.cuda import compiler_v2 as compiler
    original = compiler.load_module
    def load(path):
        path = Path(path)
        metadata = json.loads(Path(str(path)+'.json').read_text())
        if metadata['ptx_sha256'] != hashlib.sha256(path.read_bytes()).hexdigest():
            raise AssertionError('actual PTX and build sidecar disagree')
        if baseline and (metadata['arch'],metadata['flags']) not in [
                (m['arch'],m['flags']) for m in baseline['actual_build_sidecars']]:
            raise AssertionError('baseline build architecture or flags differ')
        report['actual_build_sidecars'].append(metadata)
        save()  # capture AMR metadata before production cleanup removes scratch
        return original(str(path))
    compiler.load_module = load
    try:
        yield
    finally:
        compiler.load_module = original


def preflight(repo, revision, mode):
    git = lambda *a: subprocess.check_output(['git','-C',str(repo),*a],text=True).strip()
    if git('rev-parse','HEAD') != revision or git('status','--porcelain'):
        raise ValueError('requires the exact expected clean checkout')
    source = BASE if mode == 'baseline' else PRODUCTION
    if mode == 'baseline' and revision != BASE:
        raise ValueError('baseline must be exact ed9d8c9')
    if subprocess.call(['git','-C',str(repo),'diff','--quiet',source,'--','src']):
        raise ValueError('production source differs from the frozen source revision')
    return source


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--repo',type=Path,required=True)
    parser.add_argument('--expected-revision',required=True)
    parser.add_argument('--mode',choices=('baseline','candidate'),required=True)
    parser.add_argument('--baseline',type=Path)
    parser.add_argument('--output',type=Path,required=True)
    parser.add_argument('--run-gpu',action='store_true')
    parser.add_argument('--dry-run',action='store_true')
    args = parser.parse_args()
    if args.output.exists(): parser.error('output already exists')
    selected = os.environ.get('CUDA_VISIBLE_DEVICES','')
    if not args.dry_run and (not args.run_gpu or not selected or ',' in selected or selected=='-1'):
        parser.error('--run-gpu and exactly one explicit CUDA_VISIBLE_DEVICES are required')
    repo = args.repo.resolve()
    source = preflight(repo,args.expected_revision,args.mode)
    baseline = None
    if args.mode == 'candidate':
        if args.baseline is None: parser.error('candidate requires --baseline receipt')
        baseline = json.loads(args.baseline.read_text())
    sys.path[:0] = [str(repo/'src'),str(repo),str(repo/'tests/validation/multigrid'),
                   str(repo/'tests/validation/taylor_green')]
    import CelerisLab
    if Path(CelerisLab.__file__).resolve() != repo/'src/CelerisLab/__init__.py':
        raise RuntimeError('wrong imported checkout')
    from run_case import provenance
    driver_hash = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    keys = [f'{name}:shifted={shifted}' for name in cases() for shifted in (False,True)]
    if baseline is not None: baseline_receipt(baseline,driver_hash,keys)
    report = dict(mode=args.mode,production_revision=source,source=provenance(repo),
                  driver_sha256=driver_hash,
                  device_selection=selected,completed=False,dry_run=args.dry_run,
                  operator_tolerance=[OPERATOR_ATOL,OPERATOR_RTOL],
                  rest_tolerance=[REST_ATOL,REST_RTOL],replay_tolerance=[REPLAY_ATOL,0],stops=STOPS,
                  conservation_status='diagnostic_only_no_FP32_bound',
                  actual_build_sidecars=[],default_snapshots={},default_equivalence=[],
                  projected_operator=[],union_multilevel_smoke=[],changed_remap=[])
    args.output.parent.mkdir(parents=True,exist_ok=True)
    def save():
        text = json.dumps(report,indent=2,allow_nan=False)+'\n'
        if len(text.encode()) > 8*1024*1024: raise RuntimeError('receipt exceeds 8 MiB cap')
        temporary = args.output.with_suffix(args.output.suffix+'.tmp')
        temporary.write_text(text); os.replace(temporary,args.output)
    save()
    if args.dry_run:
        print('Source preflight only; no CUDA imports or execution.'); return
    try:
        from CelerisLab.cuda.context import CudaContext
        from test_refinement_blocked_base import raw
        with CudaContext(0) as ctx, capture_builds(report,save,baseline):
            import pycuda.driver as cuda
            report['device'] = dict(name=ctx.device.name(),pci_bus_id=ctx.device.pci_bus_id(),
                                    capability=list(ctx.device.compute_capability()),
                                    driver=cuda.get_driver_version(),runtime=list(cuda.get_version()))
            nvcc = shutil.which('nvcc') or '/usr/local/cuda/bin/nvcc'
            report['nvcc_version'] = subprocess.check_output([nvcc,'--version'],text=True).strip()
            if baseline and (report['device'] != baseline.get('device')
                             or report['nvcc_version'] != baseline.get('nvcc_version')
                             or selected != baseline.get('device_selection')):
                raise ValueError('baseline GPU/device selection or toolchain differs')
            for name,case in cases().items():
                for shifted in (False,True):
                    key = f'{name}:shifted={shifted}'
                    report['active_case'] = f'default:{key}'
                    result = run_case(ctx,case,shifted,'omitted',raw)
                    report['default_snapshots'][key] = result['snapshots']
                    if baseline:
                        compare_snapshots(result['snapshots'],baseline['default_snapshots'][key])
                        explicit = run_case(ctx,case,shifted,'constant',raw)
                        compare_snapshots(explicit['snapshots'],result['snapshots'])
                        report['default_equivalence'].append(dict(case=key,passed=True))
                    save()
            if baseline:
                for name,case in cases().items():
                    for shifted in (False,True):
                        report['active_case'] = f'projected:{name}:shifted={shifted}'
                        result = run_case(ctx,case,shifted,'projected_linear',raw,True,name=='concave')
                        cfg,refs,blocks = case
                        if name == 'four':
                            variant_refs = [list(reversed(rects)) for rects in refs]
                        elif name == 'concave':
                            variant_refs = [[(4,5,7,7),(4,4,5,5)]]  # same union, disjoint decomposition
                        else:
                            variant_refs = refs
                        variant = (cfg,variant_refs,[(2,2)]*len(blocks))
                        alternate = run_case(ctx,variant,shifted,'projected_linear',raw,True)
                        compare_snapshots(alternate['snapshots'],result['snapshots'])
                        report['projected_operator'].append(dict(case=name,shifted=shifted,rows=result['operator']))
                        report['union_multilevel_smoke'].append(dict(case=name,shifted=shifted,result=result,
                                                                   alternate_bits_identical=True))
                        save()
                for shifted in (False,True):
                    report['active_case'] = f'changed-remap:shifted={shifted}'
                    report['changed_remap'].append(remap_case(ctx,shifted,raw)); save()
                expected_builds = {(s['arch'],tuple(s['flags'])) for s in baseline['actual_build_sidecars']}
                actual_builds = {(s['arch'],tuple(s['flags'])) for s in report['actual_build_sidecars']}
                if actual_builds != expected_builds:
                    raise AssertionError('baseline build architecture or compiler flags differ')
            report['completed'] = True
            save()
    except Exception as error:
        report['error'] = f'{type(error).__name__}: {error}'
        save(); raise
    print(json.dumps(dict(completed=True,mode=args.mode,receipt=str(args.output))))


if __name__ == '__main__':
    main()