package producer

import (
	"image"
	"math"

	log "github.com/sirupsen/logrus"

	"gocv.io/x/gocv"
)

func PANFilter(panImage16UC1 gocv.Mat) gocv.Mat {
	log.Println("Applying PAN filter...")
	// 转换为浮点数类型进行傅里叶变换
	imgFloat := gocv.NewMat()
	panImage16UC1.ConvertTo(&imgFloat, gocv.MatTypeCV32F)

	// 使用傅里叶变换
	planes := gocv.NewMat()
	gocv.Merge([]gocv.Mat{imgFloat, gocv.NewMatWithSize(imgFloat.Rows(), imgFloat.Cols(), gocv.MatTypeCV32F)}, &planes)
	dft := gocv.NewMat()
	gocv.DFT(planes, &dft, gocv.DftComplexOutput)

	// 转换DFT图像，使低频分量位于中心
	dftShifted := shiftDFT(dft)

	// 创建掩膜
	rows, cols := panImage16UC1.Rows(), panImage16UC1.Cols()
	crow, ccol := rows/2, cols/2
	mask := gocv.NewMatWithSize(rows, cols, gocv.MatTypeCV32FC2)
	mask.SetTo(gocv.NewScalar(1.0, 1.0, 0.0, 0.0))

	r := 30 // 调整这个参数来改变滤波效果
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			if math.Abs(float64(i-crow)) <= float64(r) && math.Abs(float64(j-ccol)) <= float64(r) {
				mask.SetFloatAt(i, j*2, 0)
				mask.SetFloatAt(i, j*2+1, 0)
			}
		}
	}

	// 应用掩膜并进行反向DFT
	filtered := gocv.NewMat()
	gocv.MulSpectrums(dftShifted, mask, &filtered, 0)
	filteredShifted := shiftDFT(filtered)
	idft := gocv.NewMat()
	gocv.DFT(filteredShifted, &idft, gocv.DftInverse|gocv.DftRealOutput)

	// 标准化图像到0-65535范围
	gocv.Normalize(idft, &idft, 0, 65535, gocv.NormMinMax)
	idft.ConvertTo(&idft, gocv.MatTypeCV16U)

	return idft
}

// shiftDFT 将DFT结果进行频谱平移
func shiftDFT(src gocv.Mat) gocv.Mat {
	rows, cols := src.Rows(), src.Cols()
	crow, ccol := rows/2, cols/2

	q0 := src.Region(image.Rect(0, 0, ccol, crow))
	q1 := src.Region(image.Rect(ccol, 0, cols, crow))
	q2 := src.Region(image.Rect(0, crow, ccol, rows))
	q3 := src.Region(image.Rect(ccol, crow, cols, rows))

	tmp := gocv.NewMatWithSize(crow, ccol, src.Type())

	q0.CopyTo(&tmp)
	q3.CopyTo(&q0)
	tmp.CopyTo(&q3)

	q1.CopyTo(&tmp)
	q2.CopyTo(&q1)
	tmp.CopyTo(&q2)

	return src
}