package sdk

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

	"github.com/gogo/protobuf/proto"

	"analysis/gohumantrack"

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

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

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

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

	t.tracker = gohumantrack.NewHumanTracker(t.gpu, t.batchSize, t.flag)
	return true
}

// Run impl
func (t *HumanTracker) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
	tm := time.Now()
	sc := 0

	for {
		select {
		case <-ctx.Done():
			return
		default:
			rMsg := <-in
			if !validRemoteMessage(rMsg, typ) {
				ejectResult(nil, rMsg, out)
				continue
			}

			i := unpackImage(rMsg, typ)
			if i == nil || i.Data == nil || i.Width <= 0 || i.Height <= 0 {
				ejectResult(nil, rMsg, out)
				continue
			}

			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)
				continue
			}

			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)

			/////////////////////////////////////
			sc++
			if sc == 25 {
				logo.Infoln("HUMAN TRACKER RUN 25 FRAME USE TIME: ", time.Since(tm))
				sc = 0
				tm = time.Now()
			}

			if time.Since(tm) > time.Second {
				logo.Infof("HUMAN TRACKER RUN %d FRAME USE TIME: %v", sc, time.Since(tm))
				sc = 0
				tm = time.Now()
			}

		}
	}
}

// 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
}