package sdk

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

	"github.com/gogo/protobuf/proto"

	"analysis/gohumantrack"

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

type imageWithID struct {
	img *gohumantrack.ImageHumanTracker
	fx  float64
	fy  float64
	id  string
}

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

	mapCameraImageIndex map[string]int
	recvImageCount      int
	index               int
	images              []*imageWithID
	msgs                []*work.MsgRS
}

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

		recvImageCount:      0,
		index:               0,
		mapCameraImageIndex: make(map[string]int),
		images:              make([]*imageWithID, batchSize),
		msgs:                make([]*work.MsgRS, batchSize),
	}
}

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

// Init impl
func (t *HumanTracker) Init() bool {

	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[0])
	result := protomsg.HumanTrackResult{Result: hr}
	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))

}

func (t *HumanTracker) trackBatch(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 imgData []byte
	rw, rh := 1280, 720
	fx, fy := 1.0, 1.0
	if imgW != rw || imgH != rh {
		imgData = goconv.ResizeBGR(i.Data, imgW, imgH, rw, rh)
		fx = (float64)(rw) / (float64)(imgW)
		fy = (float64)(rh) / (float64)(imgH)
	}
	img := gohumantrack.ImageHumanTracker{
		Data:    imgData,
		Width:   rw,
		Height:  rh,
		Channel: 3,
	}

	// mapCameraImageIndex map[string]int
	// images              []gohumantrack.ImageHumanTracker

	if i, ok := t.mapCameraImageIndex[rMsg.Msg.Cid]; ok {
		if i < t.batchSize {
			t.images[i] = &imageWithID{&img, fx, fy, rMsg.Msg.Cid}
			t.msgs[i] = &rMsg
		}
	} else {
		if t.index < t.batchSize {
			t.images[t.index] = &imageWithID{&img, fx, fy, rMsg.Msg.Cid}
			t.msgs[t.index] = &rMsg
			t.mapCameraImageIndex[rMsg.Msg.Cid] = t.index
		}
		t.index++
	}
	t.recvImageCount++

	if t.recvImageCount < t.batchSize+t.batchSize/2 {
		return
	}

	for k, v := range t.mapCameraImageIndex {
		logo.Infoln("batch~~~~~~Map index: ", v, " camera: ", k)
	}

	var pimg []*gohumantrack.ImageHumanTracker
	for k, v := range t.images[:] {
		if v != nil {
			pimg = append(pimg, v.img)
			logo.Infoln("batch~~~~~~Image index: ", k, " camera: ", v.id, " image address: ", v.img)
		} else {
			pimg = append(pimg, nil)
		}
	}

	res, err := t.tracker.ProcessImagePointer(pimg, rw, rh, 3)
	t.recvImageCount = 0

	if err != nil {
		logo.Infoln("batch~~~~~~Track Image Count: ", t.index, " Failed: ", err)
		ejectResult(nil, rMsg, out)
		return
	}

	for i := 0; i < t.batchSize; i++ {
		if t.images[i] == nil {
			continue
		}
		hr := convert2ProtoHumanTrackResultWithScale(res[i], t.images[i].fx, t.images[i].fy)
		result := protomsg.HumanTrackResult{Result: hr}
		data, err := proto.Marshal(&result)
		if err != nil {
			logo.Errorln("batch~~~~~~HUMAN TRACKER MARSHAL PROTO PLATE IDS ERROR", err)
			data = nil
		}
		msg := *t.msgs[i]
		ejectResult(data, msg, out)
		var id, name string
		if msg.Msg.Tasklab != nil {
			id, name = msg.Msg.Tasklab.Taskid, msg.Msg.Tasklab.Taskname
		}
		logo.Infoln("batch~~~~~~CAMERAID: ", msg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " Human Track COUNT: ", len(hr))

	}

}

// 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)
	FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.trackBatch, 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 {
	res := []*protomsg.HumanTrack{}
	for i := 0; i < int(obj.FgNum); i++ {
		r := obj.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)
	}
	return res
}

func convert2ProtoHumanTrackResultWithScale(obj gohumantrack.FgResult, fx, fy float64) []*protomsg.HumanTrack {
	res := []*protomsg.HumanTrack{}
	for i := 0; i < int(obj.FgNum); i++ {
		r := obj.Fginfo[i]
		rect := protomsg.Rect{
			Left:   (int32)((float64)(r.Left) / fx),
			Right:  (int32)((float64)(r.Right) / fy),
			Top:    (int32)((float64)(r.Top) / fx),
			Bottom: (int32)((float64)(r.Bottom) / fy),
		}
		pr := &protomsg.HumanTrack{
			RcHuman:    &rect,
			Confidence: r.Confidence,
			X:          (int32)((float64)(r.X) / fx),
			Y:          (int32)((float64)(r.Y) / fy),
			Id:         r.ID,
			Feature:    r.Feature[:],
		}
		res = append(res, pr)
	}
	return res
}