"""Recorded-file Pose only; independent Triton and no production network access."""
import csv
import hashlib
import json
import pathlib
import subprocess
import sys
import time
from types import SimpleNamespace
import urllib.request

import av
import cv2
import numpy as np

sys.path.insert(0,'/opt/contextra')
from app.rfdetr_postprocess import parse_boxes_labels, parse_pose_keypoints
from app.byte_tracking import ByteTracker, TrackingSettings
from app.pose_nms import pose_nms_indices

ROOT=pathlib.Path('/eval')
frames_dir=ROOT/'inputs/frames'
frames_dir.mkdir(parents=True,exist_ok=True)
config='''name: "RFDETRKeypointPreview"
platform: "tensorrt_plan"
max_batch_size: 8
version_policy { specific { versions: [8] } }
input [{ name: "input" data_type: TYPE_FP16 dims: [3,576,576] }]
output [
{ name: "dets" data_type: TYPE_FP32 dims: [100,4] },
{ name: "labels" data_type: TYPE_FP32 dims: [100,2] },
{ name: "keypoints" data_type: TYPE_FP32 dims: [100,34,8] }]
instance_group [{ kind: KIND_GPU count: 1 gpus: [0] }]
'''
(ROOT/'pose_models/RFDETRKeypointPreview/config.pbtxt').write_text(config)
rows=list(csv.DictReader(open('/work/inputs/camera24_frame_map.csv')))
pts=np.array([float(r['elapsed_seconds']) for r in rows])
grid=np.arange(400)/10.
indices=np.array([int(np.argmin(abs(pts-t))) for t in grid])
assert len(set(indices))==400 and np.all(np.diff(indices)>0)
targets={int(k+1):j for j,k in enumerate(indices)}
tracker=ByteTracker(TrackingSettings(frame_rate=10))
report=dict(status='STARTED',fps=10,sampling='nearest actual PTS to 0.1s grid; unique source frames',
    person_threshold=.15,joint_threshold=.3,nms_iou=.85,
    tracking=tracker.settings.payload(),
    preprocess='CPU OpenCV bilinear stretch RGB576, ImageNet mean/std, FP32 arithmetic -> FP16 NCHW; GPU resize bit-equivalence not claimed',
    model='RFDETRKeypointPreview',engine_sha256=hashlib.sha256((ROOT/'pose_models/RFDETRKeypointPreview/8/model.plan').read_bytes()).hexdigest(),
    production_queue_used=False,experiment_instances=1,source_files={})
for name in ('rfdetr_postprocess.py','byte_tracking.py','tracking_kernel.py','pose_nms.py'):
    src=pathlib.Path('/opt/contextra/app')/name
    dest=ROOT/'reports'/('source_'+name)
    dest.write_bytes(src.read_bytes())
    report['source_files'][name]=hashlib.sha256(src.read_bytes()).hexdigest()
log=open(ROOT/'reports/isolated_triton.log','w')
proc=subprocess.Popen(['/opt/tritonserver/bin/tritonserver','--model-repository=/eval/pose_models',
    '--model-control-mode=explicit','--load-model=RFDETRKeypointPreview','--allow-grpc=false',
    '--allow-metrics=false','--http-thread-count=2','--pinned-memory-pool-byte-size=67108864',
    '--cuda-memory-pool-byte-size=0:67108864'],stdout=log,stderr=subprocess.STDOUT)
start=time.monotonic()
all_rows=[]
try:
    for _ in range(90):
        if proc.poll() is not None: raise RuntimeError('Isolated Triton exited')
        try:
            with urllib.request.urlopen('http://127.0.0.1:8000/v2/models/RFDETRKeypointPreview/versions/8/ready',timeout=5) as r:
                if r.status==200: break
        except Exception: pass
        time.sleep(1)
    else: raise TimeoutError('Isolated Pose READY timeout')
    source=av.open('/work/inputs/camera24_original.mp4')
    source.streams.video[0].codec_context.thread_count=2
    for number,frame in enumerate(source.decode(video=0),1):
        if number>max(targets): break
        if number not in targets: continue
        sample=targets[number]
        bgr=frame.to_ndarray(format='bgr24')
        height,width=bgr.shape[:2]
        rgb=cv2.cvtColor(cv2.resize(bgr,(576,576),interpolation=cv2.INTER_LINEAR),cv2.COLOR_BGR2RGB)
        x=(rgb.astype(np.float32)/255-np.array([.485,.456,.406],np.float32))/np.array([.229,.224,.225],np.float32)
        x=np.ascontiguousarray(x.transpose(2,0,1)[None],dtype=np.float16)
        header=json.dumps(dict(inputs=[dict(name='input',shape=list(x.shape),datatype='FP16',
            parameters=dict(binary_data_size=x.nbytes))],outputs=[dict(name=n,parameters=dict(binary_data=True))
            for n in ('dets','labels','keypoints')])).encode()
        req=urllib.request.Request('http://127.0.0.1:8000/v2/models/RFDETRKeypointPreview/versions/8/infer',
            data=header+x.tobytes(),headers={'Content-Type':'application/octet-stream',
            'Inference-Header-Content-Length':str(len(header))})
        with urllib.request.urlopen(req,timeout=30) as r:
            n=int(r.headers['Inference-Header-Content-Length']); payload=r.read()
        meta=json.loads(payload[:n]); raw={}; offset=n
        for out in meta['outputs']:
            size=out['parameters']['binary_data_size']; assert out['datatype']=='FP32'
            raw[out['name']]=np.frombuffer(payload[offset:offset+size],dtype='<f4').reshape(out['shape'])
            offset+=size
        assert offset==len(payload)
        dets=raw['dets'][0]; labels=raw['labels'][0]; kp=raw['keypoints'][0]
        parsed=parse_boxes_labels(dets,labels,threshold=.1)
        kps=parse_pose_keypoints(kp,parsed['model_classes'])
        observations=[]
        for q in np.flatnonzero(parsed['keep']):
            cx,cy,bw,bh=map(float,dets[q]); box=(cx-bw/2,cy-bh/2,cx+bw/2,cy+bh/2)
            observations.append(SimpleNamespace(class_id=0,score=float(parsed['scores'][q]),bbox_xyxy=box,model_key='RFDETRKeypointPreview',
                frame=SimpleNamespace(frame_width=width,frame_height=height),query=int(q)))
        selected=pose_nms_indices(observations,.85)
        observations=[observations[i] for i in selected]
        ids=tracker.update(observations)
        persons=[]
        for o,tid in zip(observations,ids):
            if o.score < .15: continue
            points=kps[o.query].copy(); points[:,0]*=width; points[:,1]*=height
            persons.append(dict(track_id=tid,query=o.query,score=o.score,
                bbox=(np.array(o.bbox_xyxy)*[width,height,width,height]).tolist(),keypoints=points.tolist(),
                valid_joints=int((points[:,2]>=.3).sum())))
        path=frames_dir/(f'{sample:04d}.png')
        assert cv2.imwrite(str(path),bgr,[cv2.IMWRITE_PNG_COMPRESSION,1])
        all_rows.append(dict(sample=sample,target_seconds=float(grid[sample]),source_frame=number,
            seconds=float(pts[number-1]),width=width,height=height,persons=persons))
        if sample%50==0:
            print(json.dumps(dict(sample=sample,people=len(persons),wall=round(time.monotonic()-start,1))),flush=True)
            mem=int(subprocess.check_output(['nvidia-smi','--query-gpu=memory.free','--format=csv,noheader,nounits']).decode().strip())
            if mem<4096: raise RuntimeError('GPU safety margin reached')
        time.sleep(.02)
    source.close()
    assert len(all_rows)==400
    (ROOT/'inputs/pose.json').write_text(json.dumps(all_rows))
    with (ROOT/'inputs/sampling.csv').open('w') as f:
        w=csv.DictWriter(f,fieldnames=['sample','target_seconds','source_frame','seconds'])
        w.writeheader(); w.writerows({k:r[k] for k in w.fieldnames} for r in all_rows)
    report.update(status='PASS',samples=len(all_rows),people_min=min(len(r['persons']) for r in all_rows),
        people_max=max(len(r['persons']) for r in all_rows),
        untracked_observations=sum(p['track_id'] is None for r in all_rows for p in r['persons']),
        max_sampling_error_seconds=float(max(abs(pts[indices]-grid))))
finally:
    proc.terminate()
    try: proc.wait(timeout=20)
    except subprocess.TimeoutExpired: proc.kill(); proc.wait(timeout=10)
    log.close()
    report['wall_seconds']=round(time.monotonic()-start,3)
    (ROOT/'reports/pose_summary.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report),flush=True)