package yolo

import (
	"image"
	"math"

	"git.proit.id/dsutanto/bytetrack-counter-go/pkg/rknn"
	"gocv.io/x/gocv"
)

type Detection struct {
	BBox      [4]float64
	Score     float64
	Class     int
	Keypoints [][2]float64
}

type Detector struct {
	rknn         *rknn.Context
	imgSz        int
	conf         float64
	iouThr       float64
	numClasses   int
	scoreSigmoid bool
}

func NewDetector(modelPath string, imgSz int, conf float64, numClasses int, scoreSigmoid bool) (*Detector, error) {
	ctx, err := rknn.LoadModel(modelPath)
	if err != nil {
		return nil, err
	}
	return &Detector{
		rknn:         ctx,
		imgSz:        imgSz,
		conf:         conf,
		iouThr:       0.45,
		numClasses:   numClasses,
		scoreSigmoid: scoreSigmoid,
	}, nil
}

func (d *Detector) Detect(frame gocv.Mat) ([]Detection, error) {
	h0, w0 := frame.Rows(), frame.Cols()

	scale := math.Min(float64(d.imgSz)/float64(h0), float64(d.imgSz)/float64(w0))
	nh, nw := int(float64(h0)*scale), int(float64(w0)*scale)

	resized := gocv.NewMat()
	defer resized.Close()
	gocv.Resize(frame, &resized, image.Point{X: nw, Y: nh}, 0, 0, gocv.InterpolationLinear)

	letterbox := gocv.NewMatWithSize(d.imgSz, d.imgSz, gocv.MatTypeCV8UC3)
	defer letterbox.Close()
	letterbox.SetTo(gocv.NewScalar(114, 114, 114, 0))

	dy := (d.imgSz - nh) / 2
	dx := (d.imgSz - nw) / 2
	roi := letterbox.Region(image.Rect(dx, dy, dx+nw, dy+nh))
	resized.CopyTo(&roi)

	rgb := gocv.NewMat()
	defer rgb.Close()
	gocv.CvtColor(letterbox, &rgb, gocv.ColorBGRToRGB)

	data := rgb.ToBytes()

	outputs, err := d.rknn.InferenceRGB(data, d.imgSz, d.imgSz)
	if err != nil {
		return nil, err
	}
	if len(outputs) == 0 {
		return nil, nil
	}

	out := outputs[0]
	return d.decodeOutput(out, h0, w0, scale, float64(dy), float64(dx))
}

func (d *Detector) decodeOutput(out []float32, h0, w0 int, scale, padY, padX float64) ([]Detection, error) {
	if len(out) == 0 {
		return nil, nil
	}

	stride := d.numClasses + 4
	if len(out)%stride != 0 {
		return nil, nil
	}

	numDet := len(out) / stride
	if numDet == 0 {
		return nil, nil
	}

	var detections []Detection

	for i := 0; i < numDet; i++ {
		off := i * stride

		cx := float64(out[off+0])
		cy := float64(out[off+1])
		w := float64(out[off+2])
		h := float64(out[off+3])

		x1 := cx - w/2
		y1 := cy - h/2
		x2 := cx + w/2
		y2 := cy + h/2

		var maxScore float64
		var bestClass int
		for c := 0; c < d.numClasses; c++ {
			score := float64(out[off+4+c])
			if d.scoreSigmoid {
				score = sigmoid(clamp(score, -10, 10))
			}
			if score > maxScore {
				maxScore = score
				bestClass = c
			}
		}

		if maxScore <= d.conf {
			continue
		}

		x1 = (x1 - padX) / scale
		y1 = (y1 - padY) / scale
		x2 = (x2 - padX) / scale
		y2 = (y2 - padY) / scale

		x1 = clamp(x1, 0, float64(w0))
		y1 = clamp(y1, 0, float64(h0))
		x2 = clamp(x2, 0, float64(w0))
		y2 = clamp(y2, 0, float64(h0))

		detections = append(detections, Detection{
			BBox:  [4]float64{x1, y1, x2, y2},
			Score: maxScore,
			Class: bestClass,
		})
	}

	return nmsByClass(detections, d.iouThr, d.numClasses), nil
}

func nmsByClass(dets []Detection, iouThr float64, numClasses int) []Detection {
	if len(dets) == 0 {
		return nil
	}

	byClass := make([][]int, numClasses)
	for i, d := range dets {
		byClass[d.Class] = append(byClass[d.Class], i)
	}

	var result []Detection
	for cls := 0; cls < numClasses; cls++ {
		idxs := byClass[cls]
		if len(idxs) == 0 {
			continue
		}

		scores := make([]float64, len(idxs))
		boxes := make([][4]float64, len(idxs))
		for i, idx := range idxs {
			scores[i] = dets[idx].Score
			boxes[i] = dets[idx].BBox
		}

		keep := nmsIndices(boxes, scores, iouThr)
		for _, k := range keep {
			result = append(result, dets[idxs[k]])
		}
	}
	return result
}

func nmsIndices(boxes [][4]float64, scores []float64, iouThr float64) []int {
	if len(scores) == 0 {
		return nil
	}

	order := make([]int, len(scores))
	for i := range order {
		order[i] = i
	}

	for i := 0; i < len(scores); i++ {
		for j := i + 1; j < len(scores); j++ {
			if scores[order[i]] < scores[order[j]] {
				order[i], order[j] = order[j], order[i]
			}
		}
	}

	var keep []int
	for len(order) > 0 {
		idx := order[0]
		keep = append(keep, idx)
		if len(order) == 1 {
			break
		}

		rest := order[1:]
		var newOrder []int
		for _, r := range rest {
			xx1 := math.Max(boxes[idx][0], boxes[r][0])
			yy1 := math.Max(boxes[idx][1], boxes[r][1])
			xx2 := math.Min(boxes[idx][2], boxes[r][2])
			yy2 := math.Min(boxes[idx][3], boxes[r][3])
			w := math.Max(0, xx2-xx1)
			h := math.Max(0, yy2-yy1)
			inter := w * h
			areaI := (boxes[idx][2] - boxes[idx][0]) * (boxes[idx][3] - boxes[idx][1])
			areaR := (boxes[r][2] - boxes[r][0]) * (boxes[r][3] - boxes[r][1])
			iou := inter / (areaI + areaR - inter + 1e-16)
			if iou < iouThr {
				newOrder = append(newOrder, r)
			}
		}
		order = newOrder
	}
	return keep
}

func sigmoid(x float64) float64 {
	return 1.0 / (1.0 + math.Exp(-x))
}

func clamp(x, lo, hi float64) float64 {
	return math.Max(lo, math.Min(hi, x))
}

func (d *Detector) Release() {
	d.rknn.Release()
}