"""TDD for HailoPoseAdapter — locks the yolov8s_pose host decode (DFL box + keypoint
decode + NMS + COCO-17 assembly + letterbox un-mapping) WITHOUT hardware/hailort,
by injecting a fake infer_fn that returns synthetic raw output tensors.

yolov8s_pose emits 9 raw tensors (no on-chip NMS): per stride (8/16/32) a box(H,W,64
DFL), score(H,W,1 already-sigmoid), and kpts(H,W,51 = 17*(x,y,score)).
"""
import numpy as np

from edge.adapters.hailo_pose import HailoPoseAdapter

INPUT = 640
STRIDES = (8, 16, 32)


def _blank_outputs():
    outs = []
    for s in STRIDES:
        size = INPUT // s
        outs.append(np.zeros((1, size, size, 64), np.float32))  # box DFL
        outs.append(np.zeros((1, size, size, 1), np.float32))   # class score (sigmoid'd)
        outs.append(np.zeros((1, size, size, 51), np.float32))  # keypoints
    return outs


def _one_person_outputs():
    outs = _blank_outputs()
    # activate a single cell on the stride-8 map (size 80), grid cell (40, 40)
    box, score, kpts = outs[0], outs[1], outs[2]
    r = c = 40
    score[0, r, c, 0] = 0.9                       # person prob (already sigmoid on-chip)
    # DFL: one-hot each of the 4 sides at bin index 2 -> distance = 2 cells * stride8 = 16 px
    b = box[0, r, c].reshape(4, 16)
    b[:, 2] = 50.0
    box[0, r, c] = b.reshape(64)
    # keypoints: raw x=y=0.5, raw score=3.0 for all 17 joints
    k = kpts[0, r, c].reshape(17, 3)
    k[:, 0] = 0.5
    k[:, 1] = 0.5
    k[:, 2] = 3.0
    kpts[0, r, c] = k.reshape(51)
    return outs


def test_infer_maps_yolov8_pose_to_coco17_person():
    frame = np.zeros((640, 640, 3), np.uint8)   # already 640x640 -> scale 1, no padding
    outs = _one_person_outputs()
    a = HailoPoseAdapter("unused.hef", infer_fn=lambda f: outs)
    persons = a.infer(frame)
    assert len(persons) == 1
    p = persons[0]
    assert len(p.keypoints) == 17
    # grid center = (40 + 0.5) * 8 = 324 ; box +/- 16 -> [308, 340]
    assert abs(p.bbox.x1 - 308) < 2 and abs(p.bbox.y1 - 308) < 2
    assert abs(p.bbox.x2 - 340) < 2 and abs(p.bbox.y2 - 340) < 2
    # keypoint decode: stride*(2*0.5 - 0.5) + 324 = 8*0.5 + 324 = 328
    assert abs(p.keypoints[0].x - 328) < 2 and abs(p.keypoints[0].y - 328) < 2
    # joint score sigmoided: sigmoid(3.0) ~= 0.9526
    assert abs(p.keypoints[0].conf - 0.9526) < 0.01


def test_no_detection_returns_empty():
    a = HailoPoseAdapter("unused.hef", infer_fn=lambda f: _blank_outputs())
    assert a.infer(np.zeros((640, 640, 3), np.uint8)) == []