"""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-linearization-reuse-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-linearization-reuse-dense-v3/numba-cache/sparse').resolve())
os.environ['CUPY_CACHE_DIR'] = str(Path('outputs/public-glim-linearization-reuse-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_requests':0, 'initial_rank_calls':0, 'initial_rank_compiled':0, 'gpu_linearization_cache_hits':0, 'max_per_context_cache_bytes':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()
        before = self.linearization_calls
        hits = self.linearization_cache_hits
        result = native(self, arrays, *args)
        with lock:
            launched = self.linearization_calls-before if name=='linearize' else 1
            counts[label+'_calls'] += launched
            counts[label+'_rows'] += len(arrays.points)*launched
            if name=='linearize':
                counts['gpu_linearization_requests'] += 1
                counts['gpu_linearization_cache_hits'] += self.linearization_cache_hits-hits
                counts['max_per_context_cache_bytes'] = max(counts['max_per_context_cache_bytes'], self.linearization_cache_bytes)
            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):
            started = time.perf_counter()
            result = native(*args, **kwargs)
            with lock:
                counts[key] += 1
                timers[key] = timers.get(key, 0.0)+time.perf_counter()-started
            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
from calibrex.evaluation.packed_lidar_observations import _PackedObservationBuilder
for owner,name in [(_PackedObservationBuilder,'build'), (CudaPlaneQueries,'query')]:
    def observe_phase(native,label):
        def observed(*args,**kwargs):
            started = time.perf_counter()
            result = native(*args,**kwargs)
            with lock:
                timers[label] = timers.get(label,0.0)+time.perf_counter()-started
            return result
        return observed
    setattr(owner,name,observe_phase(getattr(owner,name),owner.__name__+'.'+name))
from calibrex.graph.lidar_point_to_plane import LidarRigPointToPlaneFactor
native_rank = LidarRigPointToPlaneFactor._compiled_initial_rank
def observed_rank(self, correction):
    result = native_rank(self, correction)
    with lock:
        counts['initial_rank_calls'] += 1
        counts['initial_rank_compiled'] += int(result is not None)
    return result
LidarRigPointToPlaneFactor._compiled_initial_rank = observed_rank
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-linearization-reuse-dense-v3/sparse.repeat-01/measurement.json').read_bytes())
assert b._canonical(row['result']) == b._canonical(reference['result'])
assert counts['gpu_linearization_requests'] == counts['cpu_normal_calls'] == 129
assert counts['gpu_linearization_calls']+counts['gpu_linearization_cache_hits'] == 129
assert counts['gpu_linearization_cache_hits'] > 0
assert counts['max_per_context_cache_bytes'] <= 64*1024**2
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)