package main

/*
#cgo CFLAGS: -I${SRCDIR}/darknet/include -I/usr/local/cuda-8.0/include -w -g
#cgo CXXFLAGS: -I${SRCDIR}/darknet/include -I/usr/local/cuda-8.0/include -w -g -std=c++11
#cgo LDFLAGS: -L/usr/local/cuda-8.0/lib64 -L${SRCDIR}/darknet/lib
#cgo LDFLAGS: -Wl,-rpath,${SRCDIR}/darknet/lib
#cgo LDFLAGS: -ldarknet -lcudart -lcublas -lcurand -lcudnn -lrt -ldl -lpthread
#include <stdlib.h>
#include "cyolo.h"
*/
import "C"
import "unsafe"

// CPOINT pt
type CPOINT struct {
	X int32
	Y int32
}

// CRECT rc
type CRECT struct {
	Left   int32
	Top    int32
	Right  int32
	Bottom int32
}

// CIMAGE img
type CIMAGE struct {
	Data      *uint8
	Width     int32
	Height    int32
	Channel   int32
	Pad_cgo_0 [4]byte
}

// CObjInfo yolo
type CObjInfo struct {
	RcObj CRECT
	Typ   int32
	Prob  float32
}

// CObjTrackInfo track yolo objs info
type CObjTrackInfo struct {
	ObjInfo CObjInfo
	ID      uint64
}

// YoloHandle wrap C
type YoloHandle struct {
	handle       unsafe.Pointer
	LastYoloObjs []CObjTrackInfo //yolo跟踪的上一帧信息
	LastTrackID  uint64          //yolo 被使用的ID
}

// RatioInterTrack 跟踪判断重叠阈值
const RatioInterTrack = 50 //跟踪判断重叠阈值

// SDKImage sdk image
type SDKImage struct {
	Data    []byte
	Width   int
	Height  int
	Channel int
}

// NewYolo init yolo sdk
func NewYolo(fc, fw, fn string, gi int) *YoloHandle {

	c := C.CString(fc)
	defer C.free(unsafe.Pointer(c))
	w := C.CString(fw)
	defer C.free(unsafe.Pointer(w))
	n := C.CString(fn)
	defer C.free(unsafe.Pointer(n))

	g := C.int(gi)

	p := C.init(c, w, n, g)
	return &YoloHandle{handle: p}
}

// Free free
func (y *YoloHandle) Free() {
	if y != nil {
		if y.handle != nil {
			C.release(y.handle)
		}
	}
}

// CYoloObjInfoArrayToGoArray convert cObjInfo array to go
func CYoloObjInfoArrayToGoArray(cArray unsafe.Pointer, count int) (goArray []CObjInfo) {
	p := uintptr(cArray)

	for i := 0; i < count; i++ {
		j := *(*CObjInfo)(unsafe.Pointer(p))
		goArray = append(goArray, j)
		p += unsafe.Sizeof(j)
	}
	return
}

// YoloDetect yolo detect
func (y *YoloHandle) YoloDetect(img SDKImage, thrsh float32, umns int) []CObjInfo {

	data := img.Data
	w := img.Width
	h := img.Height
	c := img.Channel
	if c == 0 {
		c = 3
	}

	var count C.int
	var cobjinfo unsafe.Pointer

	ret := C.detect(y.handle,
		unsafe.Pointer(&data[0]), C.int(w), C.int(h), C.int(c),
		C.float(thrsh), C.int(umns),
		&cobjinfo, &count)

	if ret == 0 {
		return CYoloObjInfoArrayToGoArray(unsafe.Pointer(cobjinfo), int(count))
	}
	return nil
}

// YoloObjName obj name by type
func (y *YoloHandle) YoloObjName(typ int) string {
	p := C.obj_name_by_type(y.handle, C.int(typ))

	return C.GoString(p)
}

func max(a, b int32) int32 {
	if a < b {
		return b
	}
	return a
}

func min(a, b int32) int32 {
	if a < b {
		return a
	}
	return b
}

func countInterAreaOfTwoRect(rect1 CRECT, rect2 CRECT) int32 {
	xMin := min(rect1.Left, rect2.Left)
	yMin := min(rect1.Top, rect2.Top)
	xMax := max(rect1.Right, rect2.Right)
	yMax := max(rect1.Bottom, rect2.Bottom)

	wRect1 := rect1.Right - rect1.Left
	hRect1 := rect1.Bottom - rect1.Top

	wRect2 := rect2.Right - rect2.Left
	hRect2 := rect2.Bottom - rect2.Top

	wInter := wRect1 + wRect2 - (xMax - xMin)
	hInter := hRect1 + hRect2 - (yMax - yMin)

	if (wInter <= 0) || (hInter <= 0) {
		return 0
	}

	areaInter := wInter * hInter
	areaRect1 := wRect1 * hRect1
	areaRect2 := wRect2 * hRect2
	ratio := areaInter * 100 / min(areaRect1, areaRect2)

	return ratio
}

// YoloDetectTrack yolo detect   (只识别人)
func (y *YoloHandle) YoloDetectTrack(img SDKImage, thrsh float32, umns int) (allObjs []CObjTrackInfo, newObjs []CObjTrackInfo) {

	var tmp CObjTrackInfo
	//LastYoloObjs
	detectObjs := y.YoloDetect(img, thrsh, umns)

	for _, vLast := range y.LastYoloObjs {
		for i := 0; i < len(detectObjs); i++ {
			//fmt.Println("vNew.Typ:", vNew.Typ)
			if vLast.ObjInfo.Typ == detectObjs[i].Typ { //同一类别，比如都是人体
				ratio := countInterAreaOfTwoRect(vLast.ObjInfo.RcObj, detectObjs[i].RcObj)
				if ratio >= RatioInterTrack {
					//update LastYoloObjs
					vLast.ObjInfo.RcObj = detectObjs[i].RcObj
					vLast.ObjInfo.Prob = detectObjs[i].Prob

					allObjs = append(allObjs, vLast)
					detectObjs = append(detectObjs[:i], detectObjs[i+1:]...) //从检测目标里删除已经查到的跟踪目标
					i--
					break //上一帧跟踪的目标已经找到，无需往下处理其他检测目标
				}
			}
		}
	}

	//处理新出现的目标
	if len(detectObjs) > 0 {
		for _, vAdd := range detectObjs {
			tmp.ObjInfo = vAdd
			tmp.ID = y.LastTrackID
			y.LastTrackID++

			allObjs = append(allObjs, tmp)
			newObjs = append(newObjs, tmp)
		}
	}

	//刷新上一帧的跟踪目标
	y.LastYoloObjs = allObjs

	return allObjs, newObjs
}