"""Count actual CUDA calls outside the formal timing cohort."""
import json
import os
import subprocess
import sys
import time
from argparse import Namespace
from pathlib import Path
from threading import Lock

workspace = next(p for p in Path(__file__).resolve().parents if (p/'pyproject.toml').is_file())
sys.path.insert(0, str(workspace))
os.chdir(workspace)
from tools import benchmark_adaptive_plane_search as b

out = Path('outputs/public-glim-cuda-associations-diagnostic-v3')
out.mkdir()
config = Path('examples/public_datasets/glim_versatile/online_continuous_time_compiled_dense_capture_config.yaml')
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['NUMBA_CACHE_DIR'] = str(Path('outputs/public-glim-cuda-associations-dense-v3/numba-cache/sparse').resolve())
os.environ['CUPY_CACHE_DIR'] = str(Path('outputs/public-glim-cuda-associations-dense-v3/cuda-cache/sparse').resolve())
sources = {p:b._digest(Path(p)) for p in (*b.SOURCES, config.as_posix())}
sources['diagnostic-launcher.py.txt'] = b._digest(Path(__file__))
inputs = b._verified_inputs(b.REPLAY)
revision = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()
from calibrex.solvers.cuda_continuous_time_adapter import CudaRigLinearizer
from calibrex.solvers import numba_continuous_time_adapter as a
from calibrex.solvers.cuda_plane_query_adapter import CudaPlaneQueries
from calibrex.data.adaptive_voxel import AdaptivePlaneQueryContext
from calibrex.data.livox import VoxelPlaneQueryContext
import numpy as np
from calibrex.data.numba_lidar_deskew_adapter import _PreparedDeskew
replay_contexts = []
replay_capture = _PreparedDeskew.capture
@classmethod
def capture_replay(cls, *args):
    value = replay_capture(*args)
    if value is not None:
        replay_contexts.append(value)
    return value
_PreparedDeskew.capture = capture_replay
plane_contexts = []
plane_capture = CudaPlaneQueries.capture
@classmethod
def capture_planes(cls, *args):
    value = plane_capture(*args)
    if value is not None:
        plane_contexts.append(value)
    return value
CudaPlaneQueries.capture = capture_planes
counts = {'gpu_linearization_calls':0, 'gpu_linearization_rows':0,
    'gpu_residual_calls':0, 'gpu_residual_rows':0, 'gpu_payload_uploads':0,
    'gpu_payload_bytes':0, 'cpu_linearization_calls':0, 'cpu_residual_calls':0,
    'cpu_normal_calls':0, 'cpu_uniform_query_calls':0, 'cpu_adaptive_query_calls':0, 'cpu_adaptive_repaired_rows':0}
timers = {}
lock = Lock()

def observed_gpu(name, label):
    native = getattr(CudaRigLinearizer, name)
    def observed(self, arrays, *args):
        start = time.perf_counter()
        result = native(self, arrays, *args)
        with lock:
            counts[label+'_calls'] += 1
            counts[label+'_rows'] += len(arrays.points)
            timers[label] = timers.get(label, 0.0)+time.perf_counter()-start
        return result
    setattr(CudaRigLinearizer, name, observed)

capture = CudaRigLinearizer._capture
def observed_capture(self, arrays):
    changed = self._arrays is not arrays
    result = capture(self, arrays)
    if changed:
        with lock:
            counts['gpu_payload_uploads'] += 1
            counts['gpu_payload_bytes'] += sum(v.nbytes for v in (
                arrays.points, arrays.plane_points, arrays.normals, arrays.poses, arrays.sqrt_weights))
    return result
CudaRigLinearizer._capture = observed_capture
observed_gpu('linearize', 'gpu_linearization')
observed_gpu('residuals', 'gpu_residual')
for name, key in [('numba_rig_linearization','cpu_linearization_calls'),
                  ('numba_rig_residuals','cpu_residual_calls'),
                  ('numba_rig_normal_equations','cpu_normal_calls')]:
    def wrap(native, key):
        def observed(*args, **kwargs):
            result = native(*args, **kwargs)
            with lock:
                counts[key] += 1
            return result
        return observed
    setattr(a, name, wrap(getattr(a, name), key))
for owner, key in [(VoxelPlaneQueryContext,'cpu_uniform_query_calls'),
                   (AdaptivePlaneQueryContext,'cpu_adaptive_query_calls')]:
    def wrap_query(native, key):
        def observed(*args, **kwargs):
            with lock:
                counts[key] += 1
            return native(*args, **kwargs)
        return observed
    owner._nearest_positions = wrap_query(owner._nearest_positions, key)
repair = AdaptivePlaneQueryContext._positions_from_neighbors
def observed_repair(self, points, distances, indexes, gate):
    first = distances[:,0]
    margin = 32.0*np.finfo(np.float64).eps
    radii = np.nextafter(first*(1.0+margin), np.inf)
    lower = np.nextafter(first*(1.0-margin), -np.inf)
    fast = np.isfinite(first) & (first>=1e-150) & (distances[:,1]>radii) & ((radii<=gate)|(lower>gate))
    with lock:
        counts['cpu_adaptive_repaired_rows'] += int(np.count_nonzero(~fast))
    return repair(self, points, distances, indexes, gate)
AdaptivePlaneQueryContext._positions_from_neighbors = observed_repair
b._worker(Namespace(worker='sparse', baseline_scope='compiled', candidate_query_backend='cuda',
                    config=config, output_dir=out/'worker'))
assert revision == subprocess.check_output(['git','rev-parse','HEAD'],text=True).strip()
assert sources == {p:(b._digest(Path(__file__)) if p=='diagnostic-launcher.py.txt'
                     else b._digest(Path(p))) for p in sources}
assert inputs == b._verified_inputs(b.REPLAY)
row = json.loads((out/'worker/measurement.json').read_bytes())
reference = json.loads(Path('outputs/public-glim-cuda-associations-dense-v3/sparse.repeat-01/measurement.json').read_bytes())
assert b._canonical(row['result']) == b._canonical(reference['result'])
assert counts['gpu_linearization_calls'] == 129 and counts['cpu_normal_calls'] == 129
assert counts['gpu_residual_calls'] == 0 and counts['cpu_residual_calls'] == 207
assert counts['gpu_payload_uploads'] == 26 and counts['cpu_linearization_calls'] == 0
assert len(replay_contexts) == 1
replay = replay_contexts[0]
counts.update(compiled_replay_calls=replay.calls, compiled_replay_rows=replay.rows)
assert replay.calls == 120 and replay.rows == 960_000
assert len(plane_contexts) == 1
planes = plane_contexts[0]
counts.update(gpu_plane_query_calls=planes.query_calls,
              gpu_uniform_query_rows=planes.uniform_rows,
              gpu_adaptive_query_rows=planes.adaptive_rows,
              gpu_map_upload_bytes=planes.map_upload_bytes,
              gpu_query_upload_bytes=planes.query_upload_bytes)
assert planes.query_calls == 52 and planes.uniform_rows == planes.adaptive_rows == 24_960_000
assert counts['cpu_uniform_query_calls'] == counts['cpu_adaptive_query_calls'] == 0
payload = {'provenance':{'source_commit':revision, 'source_sha256':sources,
    'input_sha256':inputs, 'measurement_sha256':b._digest(out/'worker/measurement.json')},
    'counts':counts, 'sum_seconds':timers, 'source_input_guards_passed':True,
    'complete_result_parity':True, 'limitations':[
        'Instrumented only; not pooled with formal timings. Call sums overlap across candidates.',
        'Standalone residuals, trajectory interpolation, adaptive boundary/tie repairs and ordered normal equations remain on CPU.']}
(out/'diagnostic.json').write_text(json.dumps(payload,indent=2)+'\n',encoding='utf-8',newline='\n')
print(json.dumps(payload['counts'],indent=2),flush=True)