package producer

import (
	"image"
	"math"

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

type PhaseShiftM struct {
	dx       float32
	dy       float32
	response float64
	Block    Block
}

type Block struct {
	width  int
	height int
	coord  image.Point // top-left corner of the block in the original image
}

func (r *Registrator) calculateBlockPhaseShift(panBlock, mssBlock gocv.Mat) (gocv.Point2f, float64) {
	pan := gocv.NewMat()
	mss := gocv.NewMat()

	panBlock.ConvertTo(&pan, gocv.MatTypeCV32FC1)
	defer pan.Close()
	mssBlock.ConvertTo(&mss, gocv.MatTypeCV32FC1)
	defer mss.Close()

	hann := gocv.NewMat()
	defer hann.Close()

	shift, response := gocv.PhaseCorrelate(pan, mss, hann)

	dx := shift.X
	dy := shift.Y
	log.Debugf("Block shift: dx = %f, dy = %f. response = %f", dx, dy, response)

	return shift, response
}

func (r *Registrator) DoMssPhaseShift() error {
	// 使用平均偏移量来做平移变换
	for band := 0; band < MssBands; band++ {
		var efficientShiftM int
		var xTotal, yTotal float32
		for _, shift := range r.phaseShifts[band] {
			if math.IsNaN(float64(shift.dx)) || math.IsNaN(float64(shift.dy)) {
				continue
			}
			efficientShiftM += 1
			xTotal += shift.dx
			yTotal += shift.dy
		}
		dx := xTotal / float32(efficientShiftM)
		dy := yTotal / float32(efficientShiftM)
		log.Println("Band", band+1, "average shift:", dx, dy, "efficientShiftM:", efficientShiftM)

		translationMat := gocv.NewMatWithSize(2, 3, gocv.MatTypeCV32F)
		defer translationMat.Close()

		translationMat.SetFloatAt(0, 0, 1)
		translationMat.SetFloatAt(0, 1, 0)
		translationMat.SetFloatAt(0, 2, dx)
		translationMat.SetFloatAt(1, 0, 0)
		translationMat.SetFloatAt(1, 1, 1)
		translationMat.SetFloatAt(1, 2, dy)

		alignedMss := gocv.NewMatWithSize(r.MssHeight, r.MssWidth, gocv.MatTypeCV32FC1)
		defer alignedMss.Close()

		cvtMss := gocv.NewMat()
		defer cvtMss.Close()

		r.MssImages[band].ConvertTo(&cvtMss, gocv.MatTypeCV32FC1)
		// 手动平移像素
		outY := math.MaxInt
		for y := 0; y < r.MssHeight; y++ {
			for x := 0; x < r.MssWidth; x++ {
				// 计算新的坐标
				newX := x + int(dx)
				newY := y + int(dy)

				// 如果新坐标在图像范围内，进行像素值赋值
				if newX >= 0 && newX < r.MssWidth && newY >= 0 && newY < r.MssHeight {
					alignedMss.SetFloatAt(y, x, cvtMss.GetFloatAt(newY, newX))
				} else {
					// 如果新坐标不在图像范围内，设置为黑色
					alignedMss.SetFloatAt(y, x, 0)
					if outY > y {
						outY = y
						log.Println("Warning: pixel out of range", x, y)
					}
				}
			}
		}

		// gocv.WarpAffine(cvtMss, &alignedMss, translationMat, image.Pt(cvtMss.Size()[1], cvtMss.Size()[0]))
		r.registeredMssImages[band] = gocv.NewMat()
		alignedMss.ConvertTo(&r.registeredMssImages[band], gocv.MatTypeCV16U)

		log.Println("Band", band+1, "registeredMssImages size:", r.registeredMssImages[band].Size())
	}

	return nil
}