package sdk

import (
	"analysis/logo"
	"analysis/work"
	"context"

	"github.com/gogo/protobuf/proto"

	"analysis/gohumantrack"

	"basic.com/pubsub/protomsg.git"
	"basic.com/valib/gogpu.git"
)

// HumanTracker track
type HumanTracker struct {
	tracker   *gohumantrack.HumanTracker
	gpu       int
	batchSize int
	flag      int
	list      *LockList
}

// NewHumanTracker new
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
	return &HumanTracker{
		gpu:       gpu,
		batchSize: batchSize,
		flag:      flag,
		list:      NewLockList(6),
	}
}

// Free free
func (t *HumanTracker) Free() {
	if t.tracker != nil {
		t.tracker.Free()
	}
}

// Init impl
func (t *HumanTracker) Init() bool {
	if t.batchSize != 1 {
		logo.Errorln("ONLY SUPPORT BATCH SIZE = 1")
		return false
	}
	gpu := t.gpu

	if gpu == -1 {
		gpu = gogpu.ValidGPU(2048)
	}
	logo.Infof("start gpu %d, batch %d, flag %d\n", gpu, t.batchSize, t.flag)
	t.tracker = gohumantrack.NewHumanTracker(gpu, t.batchSize, t.flag)
	if t.tracker != nil {
		logo.Infoln("Start Success")
		return true
	}
	logo.Infoln("Start Failed")
	return false
}

func (t *HumanTracker) track(rMsg work.MsgRS, out chan<- work.MsgRS, typ string) {
	i := unpackImage(rMsg, typ)
	if i == nil || i.Data == nil || i.Width <= 0 || i.Height <= 0 {
		ejectResult(nil, rMsg, out)
		return
	}

	imgW, imgH := int(i.Width), int(i.Height)

	var images []gohumantrack.ImageHumanTracker
	img := gohumantrack.ImageHumanTracker{
		Data:    i.Data,
		Width:   imgW,
		Height:  imgH,
		Channel: 3,
	}
	images = append(images, img)
	res, err := t.tracker.Process(images)
	if err != nil {
		ejectResult(nil, rMsg, out)
		return
	}

	hr := convert2ProtoHumanTrackResult(res)
	result := protomsg.HumanTrackResult{Result: hr[0]}
	data, err := proto.Marshal(&result)
	if err != nil {
		logo.Errorln("HUMAN TRACKER MARSHAL PROTO PLATE IDS ERROR", err)
		data = nil
	}
	ejectResult(data, rMsg, out)
	var id, name string
	if rMsg.Msg.Tasklab != nil {
		id, name = rMsg.Msg.Tasklab.Taskid, rMsg.Msg.Tasklab.Taskname
	}
	logo.Infoln("CAMERAID: ", rMsg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " Human Track COUNT: ", len(hr[0]))

}

// Run impl
func (t *HumanTracker) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
	FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.track, t.Free)
}

// message HumanTrack {
//     Rect rcHuman = 1;
// 	float confidence = 2;
// 	int32 x = 3;
// 	int32 y = 4;
// 	int32 id = 5;
// 	repeated float feature = 6;
// }

// message HumanTrackResult {
//     repeated HumanTrack result = 1;
// }

func convert2ProtoHumanTrackResult(obj []gohumantrack.FgResult) [][]*protomsg.HumanTrack {
	ret := [][]*protomsg.HumanTrack{}
	for _, v := range obj {
		res := []*protomsg.HumanTrack{}
		for i := 0; i < int(v.FgNum); i++ {
			r := v.Fginfo[i]

			rect := protomsg.Rect{
				Left:   r.Left,
				Right:  r.Right,
				Top:    r.Top,
				Bottom: r.Bottom,
			}
			pr := &protomsg.HumanTrack{
				RcHuman:    &rect,
				Confidence: r.Confidence,
				X:          r.X,
				Y:          r.Y,
				Id:         r.ID,
				Feature:    r.Feature[:],
			}
			res = append(res, pr)
		}
		ret = append(ret, res)
	}
	return ret
}