package iou

import (
	"math"
	"sort"
)

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

	keep := make([]int, 0, n)
	for len(idx) > 0 {
		cur := idx[0]
		keep = append(keep, cur)
		if len(idx) == 1 {
			break
		}
		rest := idx[1:]
		newOrder := make([]int, 0, len(rest))
		bCur := boxes[cur]
		areaCur := (bCur[2] - bCur[0]) * (bCur[3] - bCur[1])
		for _, ri := range rest {
			bRi := boxes[ri]
			xx1 := math.Max(bCur[0], bRi[0])
			yy1 := math.Max(bCur[1], bRi[1])
			xx2 := math.Min(bCur[2], bRi[2])
			yy2 := math.Min(bCur[3], bRi[3])
			w := math.Max(0, xx2-xx1)
			h := math.Max(0, yy2-yy1)
			inter := w * h
			areaRi := (bRi[2] - bRi[0]) * (bRi[3] - bRi[1])
			ov := inter / (areaCur + areaRi - inter + 1e-16)
			if ov < iouThr {
				newOrder = append(newOrder, ri)
			}
		}
		idx = newOrder
	}
	return keep
}

func PairwiseIoU(boxesA, boxesB [][4]float64) [][]float64 {
	n, m := len(boxesA), len(boxesB)
	mat := make([][]float64, n)
	for i := range mat {
		mat[i] = make([]float64, m)
	}
	if n == 0 || m == 0 {
		return mat
	}
	for i := 0; i < n; i++ {
		areaA := (boxesA[i][2] - boxesA[i][0]) * (boxesA[i][3] - boxesA[i][1])
		for j := 0; j < m; j++ {
			xx1 := math.Max(boxesA[i][0], boxesB[j][0])
			yy1 := math.Max(boxesA[i][1], boxesB[j][1])
			xx2 := math.Min(boxesA[i][2], boxesB[j][2])
			yy2 := math.Min(boxesA[i][3], boxesB[j][3])
			iw := math.Max(0, xx2-xx1)
			ih := math.Max(0, yy2-yy1)
			inter := iw * ih
			areaB := (boxesB[j][2] - boxesB[j][0]) * (boxesB[j][3] - boxesB[j][1])
			mat[i][j] = inter / (areaA + areaB - inter + 1e-16)
		}
	}
	return mat
}

func GreedyMatchRowCol(costMatrix [][]float64, threshold float64) []struct{ Row, Col int } {
	if len(costMatrix) == 0 || len(costMatrix[0]) == 0 {
		return nil
	}
	n, m := len(costMatrix), len(costMatrix[0])

	type entry struct {
		cost     float64
		row, col int
	}
	entries := make([]entry, 0, n*m)
	for i := 0; i < n; i++ {
		for j := 0; j < m; j++ {
			entries = append(entries, entry{costMatrix[i][j], i, j})
		}
	}
	sort.Slice(entries, func(a, b int) bool {
		return entries[a].cost < entries[b].cost
	})

	rowUsed := make(map[int]bool)
	colUsed := make(map[int]bool)
	var pairs []struct{ Row, Col int }
	for _, e := range entries {
		if e.cost >= threshold {
			break
		}
		if rowUsed[e.row] || colUsed[e.col] {
			continue
		}
		rowUsed[e.row] = true
		colUsed[e.col] = true
		pairs = append(pairs, struct{ Row, Col int }{e.row, e.col})
	}
	return pairs
}