package rrc

import (
	"image"
	"math"

	"gocv.io/x/gocv"
	"starwiz.cn/sjy01/image-proc/pkg/config"
)

// filter high frequence noise
func HFNoiseFilter(img gocv.Mat, radius float64) gocv.Mat {
	minGrayVal, maxGrayVal, _, _ := gocv.MinMaxLoc(img)

	// 将图像转换为32位浮点型
	img32f := gocv.NewMat()
	img.ConvertTo(&img32f, gocv.MatTypeCV32F)
	defer img32f.Close()

	// 扩展图像尺寸，使其宽高为2的幂次方
	// optRows := gocv.GetOptimalDFTSize(img32f.Rows())
	// optCols := gocv.GetOptimalDFTSize(img32f.Cols())
	// padded := gocv.NewMatWithSize(optRows, optCols, gocv.MatTypeCV32F)
	// defer padded.Close()
	// gocv.CopyMakeBorder(img32f, &padded,
	// 	0, optRows-img32f.Rows(), 0, optCols-img32f.Cols(),
	// 	gocv.BorderConstant, color.RGBA{0, 0, 0, 0})

	// 获取图像尺寸
	rows, cols := img32f.Rows(), img32f.Cols()

	// 执行 DFT
	DFT := gocv.NewMat()
	defer DFT.Close()
	zero := gocv.Zeros(rows, cols, gocv.MatTypeCV32F)
	defer zero.Close()
	gocv.Merge([]gocv.Mat{img32f, zero}, &DFT)
	gocv.DFT(DFT, &DFT, gocv.DftComplexOutput)

	// 中心化频谱
	fftshift(&DFT)

	// visualizeDFT(DFT, "/Users/lan/workspace/sjy01/image-proc/data/fft/dft.jpg")

	// 低通滤波器
	pss := gocv.Split(DFT)
	defer pss[0].Close()
	defer pss[1].Close()
	filter := gocv.NewMatWithSize(rows, cols, gocv.MatTypeCV32F)
	defer filter.Close()
	// createLowPassFilter(&filter, radius)
	// createBandBlockFilter(&filter, radius, radius+0.015*float64(cols))
	createFixedBlockFilter(&filter, config.GCONFIG.Radiation.HFBandStopWidth)
	gocv.Multiply(pss[0], filter, &pss[0])
	gocv.Multiply(pss[1], filter, &pss[1])
	gocv.Merge(pss, &DFT)

	// visualizeDFT(DFT, "/Users/lan/workspace/sjy01/image-proc/data/fft/dft_filtered.jpg")

	// 频谱位置还原
	fftshift(&DFT)

	// 执行 IDFT
	IDFT := gocv.NewMat()
	gocv.DFT(DFT, &IDFT, gocv.DftInverse|gocv.DftRealOutput)

	// 归一化到原始图像范围
	gocv.Normalize(IDFT, &IDFT, float64(minGrayVal), float64(maxGrayVal), gocv.NormMinMax)
	IDFT.ConvertTo(&IDFT, gocv.MatTypeCV16U)
	return IDFT
}

// 创建低通滤波器
func createLowPassFilter(filter *gocv.Mat, radius float64) {
	rows, cols := filter.Rows(), filter.Cols()
	center := image.Pt(cols/2, rows/2)
	for y := 0; y < rows; y++ {
		for x := 0; x < cols; x++ {
			dist := math.Sqrt(math.Pow(float64(x-center.X), 2) + math.Pow(float64(y-center.Y), 2))
			if dist <= float64(radius) {
				filter.SetFloatAt(y, x, 1.0)
			} else {
				filter.SetFloatAt(y, x, 0.0)
			}
		}
	}
}

// 创建带阻滤波器
func createBandBlockFilter(filter *gocv.Mat, radius0, radius1 float64) {
	rows, cols := filter.Rows(), filter.Cols()
	center := image.Pt(cols/2, rows/2)
	for y := 0; y < rows; y++ {
		for x := 0; x < cols; x++ {
			dist := math.Sqrt(math.Pow(float64(x-center.X), 2) + math.Pow(float64(y-center.Y), 2))
			if dist <= float64(radius0) || dist >= float64(radius1) {
				filter.SetFloatAt(y, x, 1.0)
			} else {
				filter.SetFloatAt(y, x, 0.0)
			}
		}
	}
}

func createFixedBlockFilter(filter *gocv.Mat, ksize int) {
	rows, cols := filter.Rows(), filter.Cols()
	center := image.Pt(cols/2, rows/2)
	for y := 0; y < rows; y++ {
		for x := 0; x < cols; x++ {
			if (x >= center.X-ksize/2 && x <= center.X+ksize/2 &&
				(y < int(ksize) || y >= rows-int(ksize))) ||
				((x < ksize || x > cols-ksize) && (y >= center.Y-ksize/2 && y <= center.Y+ksize/2)) {
				filter.SetFloatAt(y, x, 0.0)
			} else {
				filter.SetFloatAt(y, x, 1.0)
			}
		}
	}
}

// fftshift 将频率直流分量移动到图像中心
func fftshift(input *gocv.Mat) {
	cx, cy := input.Cols()/2, input.Rows()/2

	q0 := input.Region(image.Rect(0, 0, cx, cy))                       // Top-Left
	q1 := input.Region(image.Rect(cx, 0, input.Cols(), cy))            // Top-Right
	q2 := input.Region(image.Rect(0, cy, cx, input.Rows()))            // Bottom-Left
	q3 := input.Region(image.Rect(cx, cy, input.Cols(), input.Rows())) // Bottom-Right

	tmp := gocv.NewMatWithSize(q0.Rows(), q0.Cols(), input.Type())
	defer tmp.Close()

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

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

	q0.Close()
	q1.Close()
	q2.Close()
	q3.Close()
}

func visualizeDFT(DFT gocv.Mat, file string) {
	pss := gocv.Split(DFT)
	planes := gocv.NewMat()
	defer planes.Close()

	gocv.Magnitude(pss[0], pss[1], &planes)
	// 进行对数尺度变换，便于可视化
	ones := gocv.NewMatFromScalar(gocv.NewScalar(1, 0, 0, 0), gocv.MatTypeCV64F)
	defer ones.Close()
	gocv.Add(planes, ones, &planes)
	gocv.Log(planes, &planes)

	// 归一化
	gocv.Normalize(planes, &planes, 0, 255, gocv.NormMinMax)
	planes.ConvertTo(&planes, gocv.MatTypeCV8U)
	out := gocv.NewMatWithSize(planes.Rows()/4, planes.Cols()/4, gocv.MatTypeCV8U)
	defer out.Close()
	gocv.Resize(planes, &out, image.Point{planes.Cols() / 4, planes.Rows() / 4}, 0, 0, gocv.InterpolationLinear)

	// 保存DFT结果
	gocv.IMWrite(file, out)
}
