"""Read-only production-kernel parity diagnostic; no altered math or thresholds.
Run in checkout with PYTHONPATH=src. Supply the exact failed subcase.
"""
import argparse
import importlib.util
import json
from pathlib import Path
import shutil
import sys
import numpy as np

p = argparse.ArgumentParser()
p.add_argument('--shifted', type=int, choices=(0, 1), required=True)
p.add_argument('--profile', choices=('uniform', 'parabolic'), required=True)
p.add_argument('--clamp', type=int, choices=(0, 1), required=True)
p.add_argument('--damping', type=float, required=True)
p.add_argument('--block', type=int, nargs=2, required=True, metavar=('BY', 'BX'))
p.add_argument('--steps', type=int, default=8)
p.add_argument('--output', type=Path, required=True)
a = p.parse_args()
a.output.mkdir(parents=True, exist_ok=False)
sys.path.insert(0, str(Path('src').resolve()))
path = Path('tests/validation/multigrid/test_channel_root_gpu.py').resolve()
spec = importlib.util.spec_from_file_location('channel_gpu_gate', path)
gate = importlib.util.module_from_spec(spec); spec.loader.exec_module(gate)

# Preserve the exact NVCC-generated PTX before the existing compiler cleans up.
original_compile = gate.compiler.compile_kernel
original_module = gate.compile_module
cache = {}
def retain_compile(*args, **kwargs):
    result = original_compile(*args, **kwargs)
    source = kwargs.get('source', 'uniform.cu')
    shutil.copyfile(result, a.output / (Path(source).stem + '.ptx'))
    metadata = Path(result + '.json')
    if metadata.exists(): shutil.copyfile(metadata, a.output / (Path(source).stem + '.ptx.json'))
    return result
def cached_module(cfg, context, source):
    cfg.validate()
    key = (repr(cfg), source)
    if key not in cache: cache[key] = original_module(cfg, context, source)
    return cache[key]
gate.compiler.compile_kernel = retain_compile
gate.compile_module = cached_module

cfg = gate.LBMConfig(nx=24, ny=20, viscosity=.1, velocity=.025,
                    ddf_shifting=bool(a.shifted), inlet_profile=a.profile,
                    outlet_backflow_clamp=bool(a.clamp), outlet_srt_neq_damp=a.damping)
config = dict(shifted=a.shifted, profile=a.profile, clamp=a.clamp,
              damping=a.damping, block=a.block, steps=a.steps)
(a.output / 'config.json').write_text(json.dumps(config, indent=2))
def ordered(bits):
    return int((~bits & 0xffffffff) if bits & 0x80000000 else (bits | 0x80000000))
def report(expected, actual, step, label):
    ub = expected.view(np.uint32); ab = actual.view(np.uint32)
    indices = np.argwhere(ub != ab)
    numeric = ~((expected == actual) | (np.isnan(expected) & np.isnan(actual)))
    numeric_count = int(np.count_nonzero(numeric))
    rows = []
    for q, y, x in indices:
        u, b = int(ub[q,y,x]), int(ab[q,y,x])
        category = ('wall_ghost' if y in (0,cfg.ny-1) else 'inlet_ghost' if x==0
                    else 'outlet_ghost' if x==cfg.nx-1 else 'wall_adjacent_fluid' if y in (1,cfg.ny-2) else 'bulk_fluid')
        rows.append(dict(numeric_mismatch=bool(numeric[q,y,x]), q=int(q), y=int(y), x=int(x), category=category,
                         uniform=float(expected[q,y,x]), blocked=float(actual[q,y,x]),
                         uniform_hex=f'{u:08x}', blocked_hex=f'{b:08x}',
                         ulps=abs(ordered(u)-ordered(b))))
    record = dict(label=label, step=step, bit_count=len(rows), numeric_count=numeric_count,
                  max_abs=float(np.max(np.abs(actual.astype(float)-expected.astype(float)))),
                  max_ulps=max((r['ulps'] for r in rows), default=0),
                  all_finite=bool(np.isfinite(expected).all() and np.isfinite(actual).all()),
                  differences=rows)
    (a.output / f'{label}_{step:03d}.json').write_text(json.dumps(record,indent=2))
    print(json.dumps({k:v for k,v in record.items() if k!='differences'}), flush=True)
    if rows: print(json.dumps(rows[:12]), flush=True)
    return numeric_count > 0

with gate.CudaContext(0) as context:
    initial = gate._encode_physical(gate.state(cfg.ny,cfg.nx), bool(a.shifted))
    uniform, blocked = initial.copy(), initial.copy()
    found = False
    for step in range(1,a.steps+1):
        previous_uniform, previous_blocked = uniform, blocked
        uniform = gate.run_root(cfg, context, previous_uniform, steps=1)
        blocked = gate.run_root(cfg, context, previous_blocked, tuple(a.block), steps=1)
        differs = report(uniform, blocked, step, 'propagated')
        if differs and not found:
            found = True
            same_input_blocked = gate.run_root(cfg, context, previous_uniform, tuple(a.block), steps=1)
            report(uniform, same_input_blocked, step, 'same_input')
            np.savez(a.output/'first_divergence.npz', initial=initial,
                     previous_uniform=previous_uniform, previous_blocked=previous_blocked,
                     uniform=uniform, blocked=blocked, same_input_blocked=same_input_blocked)
    # Verify repeated one-step diagnostics still reproduce the original eight-hop call.
    direct_uniform = gate.run_root(cfg, context, initial, steps=a.steps)
    direct_blocked = gate.run_root(cfg, context, initial, tuple(a.block), steps=a.steps)
    reproduced = report(direct_uniform, direct_blocked, a.steps, 'original_multistep')
    report(direct_uniform, uniform, a.steps, 'uniform_replay_consistency')
    report(direct_blocked, blocked, a.steps, 'blocked_replay_consistency')

(a.output / 'result.json').write_text(json.dumps(dict(numeric_failure_reproduced=reproduced, first_numeric_divergence_found=found), indent=2))
print(json.dumps(dict(numeric_failure_reproduced=reproduced)), flush=True)