package pose

import "testing"

func TestParseYOLOv8PoseRestoresCoordinatesAndSuppressesOverlap(t *testing.T) {
	output := make([]float32, YOLOPoseOutputValues)
	putDetection(output, 0, 320, 320, 100, 200, 0.90, 330, 340)
	putDetection(output, 1, 322, 321, 100, 200, 0.80, 332, 341)

	people, err := ParseYOLOv8Pose(output, LetterboxTransform{Scale: 1}, 0.25, 0.70)
	if err != nil {
		t.Fatalf("ParseYOLOv8Pose returned an error: %v", err)
	}
	if len(people) != 1 {
		t.Fatalf("people = %d, want one NMS survivor", len(people))
	}
	person := people[0]
	if person.Box != (Box{Left: 270, Top: 220, Right: 370, Bottom: 420}) {
		t.Fatalf("box = %+v", person.Box)
	}
	if point := person.Keypoints[0]; point.X != 330 || point.Y != 340 || point.Confidence != 0.9 {
		t.Fatalf("first keypoint = %+v", point)
	}
}

func TestParseYOLOv8PoseRejectsUnexpectedOutputLength(t *testing.T) {
	_, err := ParseYOLOv8Pose([]float32{0}, LetterboxTransform{Scale: 1}, 0.25, 0.70)
	if err == nil {
		t.Fatal("ParseYOLOv8Pose accepted a malformed output")
	}
}

func putDetection(output []float32, candidate int, centerX, centerY, width, height, confidence, keypointX, keypointY float32) {
	output[0*YOLOPoseCandidateCount+candidate] = centerX
	output[1*YOLOPoseCandidateCount+candidate] = centerY
	output[2*YOLOPoseCandidateCount+candidate] = width
	output[3*YOLOPoseCandidateCount+candidate] = height
	output[4*YOLOPoseCandidateCount+candidate] = confidence
	for keypoint := 0; keypoint < YOLOPoseKeypointCount; keypoint++ {
		base := 5 + keypoint*3
		output[(base+0)*YOLOPoseCandidateCount+candidate] = keypointX
		output[(base+1)*YOLOPoseCandidateCount+candidate] = keypointY
		output[(base+2)*YOLOPoseCandidateCount+candidate] = 0.9
	}
}