package producer

import (
	"image"

	log "github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
)

// align the image via Enhanced Correlation Coefficient (ECC) algorithm

func CV_ECCAlign(templateImage, inputImage gocv.Mat) gocv.Mat {
	retMatType := templateImage.Type()
	warpMode := gocv.MotionHomography
	var warpMatrix gocv.Mat
	switch warpMode {
	case gocv.MotionHomography:
		warpMatrix = gocv.NewMatWithSize(3, 3, gocv.MatTypeCV32FC1)
	default:
		warpMatrix = gocv.NewMatWithSize(2, 3, gocv.MatTypeCV32FC1)
	}
	for i := 0; i < warpMatrix.Rows(); i++ {
		warpMatrix.SetFloatAt(i, i, 1.0)
	}

	criteria := gocv.NewTermCriteria(gocv.Count+gocv.EPS, 1000, 0.0001)

	var mask = gocv.NewMat()
	templateImage.ConvertTo(&templateImage, gocv.MatTypeCV32FC1)
	inputImage.ConvertTo(&inputImage, gocv.MatTypeCV32FC1)

	retval := gocv.FindTransformECC(templateImage, inputImage, &warpMatrix, warpMode, criteria, mask, 5)
	log.Info("cv::findTransformECC ret:", retval)

	var aligned = gocv.NewMat()
	sz := image.Point{inputImage.Cols(), inputImage.Rows()}
	if warpMode == gocv.MotionHomography {
		gocv.WarpPerspective(inputImage, &aligned, warpMatrix, sz)
	} else {
		gocv.WarpAffine(inputImage, &aligned, warpMatrix, sz)
	}

	aligned.ConvertTo(&aligned, retMatType)
	templateImage.ConvertTo(&templateImage, retMatType)
	inputImage.ConvertTo(&inputImage, retMatType)
	return aligned
}

func (r *ImgProc) DoECCAlign() error {
	r.registeredMssImages[0] = r.MssImages[0].Clone()
	for i := 1; i < len(r.MssImages); i++ {
		r.registeredMssImages[i] = CV_ECCAlign(r.registeredMssImages[0], r.MssImages[i])
	}
	return nil
}