package pose

import (
	"fmt"
	"math"
)

const letterboxPadding = byte(114)

// PreprocessBGR reproduces the fixed-shape Ultralytics letterbox contract:
// BGR source pixels are resized with bilinear interpolation, padded in 114
// gray, converted to RGB/CHW and normalised to [0, 1].
func PreprocessBGR(frame []byte, width, height, target int) ([]float32, LetterboxTransform, error) {
	if width <= 0 || height <= 0 || target <= 0 {
		return nil, LetterboxTransform{}, fmt.Errorf("frame width, height and target must be positive")
	}
	if len(frame) != width*height*3 {
		return nil, LetterboxTransform{}, fmt.Errorf("BGR frame length = %d, want %d", len(frame), width*height*3)
	}

	scale := math.Min(float64(target)/float64(width), float64(target)/float64(height))
	resizedWidth := int(math.Round(float64(width) * scale))
	resizedHeight := int(math.Round(float64(height) * scale))
	padWidth := target - resizedWidth
	padHeight := target - resizedHeight
	padLeft := int(math.Round(float64(padWidth)/2.0 - 0.1))
	padTop := int(math.Round(float64(padHeight)/2.0 - 0.1))
	transform := LetterboxTransform{
		OriginalWidth: width, OriginalHeight: height, InputSize: target,
		ResizedWidth: resizedWidth, ResizedHeight: resizedHeight,
		PadLeft: padLeft, PadTop: padTop, Scale: float32(scale),
	}

	plane := target * target
	input := make([]float32, 3*plane)
	padding := float32(letterboxPadding) / 255.0
	for index := range input {
		input[index] = padding
	}
	for y := 0; y < resizedHeight; y++ {
		sourceY := clampFloat((float64(y)+0.5)*float64(height)/float64(resizedHeight)-0.5, 0, float64(height-1))
		y0 := int(math.Floor(sourceY))
		y1 := minInt(y0+1, height-1)
		yWeight := float32(sourceY - float64(y0))
		for x := 0; x < resizedWidth; x++ {
			sourceX := clampFloat((float64(x)+0.5)*float64(width)/float64(resizedWidth)-0.5, 0, float64(width-1))
			x0 := int(math.Floor(sourceX))
			x1 := minInt(x0+1, width-1)
			xWeight := float32(sourceX - float64(x0))
			leftTop := (y0*width + x0) * 3
			rightTop := (y0*width + x1) * 3
			leftBottom := (y1*width + x0) * 3
			rightBottom := (y1*width + x1) * 3
			destination := (padTop+y)*target + padLeft + x
			for sourceChannel := 0; sourceChannel < 3; sourceChannel++ {
				value := bilinear(
					frame[leftTop+sourceChannel], frame[rightTop+sourceChannel],
					frame[leftBottom+sourceChannel], frame[rightBottom+sourceChannel],
					xWeight, yWeight,
				)
				// BGR source maps to RGB tensor planes.
				tensorChannel := 2 - sourceChannel
				// cv2.resize writes uint8 pixels before Ultralytics converts the
				// image to float; round here to retain that observable contract.
				input[tensorChannel*plane+destination] = float32(math.Round(float64(value))) / 255.0
			}
		}
	}
	return input, transform, nil
}

func bilinear(topLeft, topRight, bottomLeft, bottomRight byte, xWeight, yWeight float32) float32 {
	top := float32(topLeft)*(1-xWeight) + float32(topRight)*xWeight
	bottom := float32(bottomLeft)*(1-xWeight) + float32(bottomRight)*xWeight
	return top*(1-yWeight) + bottom*yWeight
}

func clampFloat(value, minimum, maximum float64) float64 {
	return math.Max(minimum, math.Min(maximum, value))
}

func minInt(left, right int) int {
	if left < right {
		return left
	}
	return right
}