import ctypes
import json
import os
from pathlib import Path
import numpy as np
import torch
from rfdetr.models.postprocess import PostProcess

os.umask(0o077)
R=Path('/evidence');report={}
torch.set_num_threads(2)
for name,shape,atol in [('dets',(1,300,4),.005),('labels',(1,300,3),.05)]:
    actual=np.fromfile(R/(name+'.bin'),dtype=np.float32).reshape(shape)
    expected=np.load(R/(name+'_onnx.npy'),allow_pickle=False)
    assert actual.shape==expected.shape and np.isfinite(actual).all()
    report[name]=dict(shape=list(shape),finite=True,atol=atol,rtol=.01,
                      max_abs=float(np.abs(actual-expected).max()),
                      numerical_pass=bool(np.allclose(actual,expected,atol=atol,rtol=.01)))
boxes=np.load(R/'dets_onnx.npy',allow_pickle=False)
logits=np.load(R/'labels_onnx.npy',allow_pickle=False)
out=PostProcess(num_select=300)({'pred_boxes':torch.from_numpy(boxes),
                              'pred_logits':torch.from_numpy(logits)},torch.tensor([[720,1280]]))[0]
indices=out['labels'].numpy()<2
labels=out['labels'].numpy()[indices]
scores=out['scores'].numpy()[indices]
xyxy=out['boxes'].numpy()[indices].copy()
xyxy[:,0::2]=xyxy[:,0::2].clip(0,1280);xyxy[:,1::2]=xyxy[:,1::2].clip(0,720)
valid=(xyxy[:,2]>xyxy[:,0])&(xyxy[:,3]>xyxy[:,1])
expected=np.column_stack((labels,scores,xyxy[:,0:2],xyxy[:,2:4]-xyxy[:,0:2]))[valid]
lib=ctypes.CDLL('/evidence/fire-decode.so')
fn=lib.fire_decode;fn.argtypes=[ctypes.c_void_p]*3;fn.restype=ctypes.c_int
actual=np.zeros((300,6),np.float32)
count=fn(boxes.ctypes.data,logits.ctypes.data,actual.ctypes.data)
actual=actual[:count]
assert actual.shape==expected.shape
assert np.array_equal(actual[:,0],expected[:,0])
assert np.allclose(actual[:,1:],expected[:,1:],atol=.0002,rtol=.00001)
report['decoder_oracle']=dict(status='PASS',count=count,max_abs=float(np.abs(actual-expected).max()),
    scope='official PostProcess top300 then foreground0/1 and valid clipped source box; no NMS')
(R/'comparison.json').write_text(json.dumps(report,indent=2))
print(json.dumps(report,indent=2))