package nms

import (
	"math"
	"sort"
)

func Apply(boxes [][4]float64, scores []float64, iouThr float64) []int {
	n := len(scores)
	order := make([]int, n)
	for i := range order {
		order[i] = i
	}
	sort.Slice(order, func(i, j int) bool {
		return scores[order[i]] > scores[order[j]]
	})

	keep := make([]int, 0)
	for len(order) > 0 {
		idx := order[0]
		keep = append(keep, idx)
		if len(order) == 1 {
			break
		}
		rest := order[1:]
		newOrder := make([]int, 0, len(rest))
		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
}