package rrc

import (
	"math"

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

// Destriping multisensor imagery with moment matching [Gadallah, 2000]
func DoMomentMatching(originalImg gocv.Mat) {
	probes := originalImg.Cols()
	log.Printf("do moment matching for %d probes", probes)
	// 第i个探元的像元均值和标准差
	means := make([]float64, probes)
	stds := make([]float64, probes)
	// 计算每个探元的均值和标准差
	for x := 0; x < originalImg.Cols(); x++ {
		var total int64
		n := originalImg.Rows()
		var dn uint16
		for y := 0; y < n; y++ {
			dn = uint16(originalImg.GetShortAt(y, x))
			total += int64(dn)
		}
		means[x] = float64(total) / float64(n)

		var a float64
		for y := 0; y < n; y++ {
			dn = uint16(originalImg.GetShortAt(y, x))
			a += math.Pow(float64(dn)-means[x], 2)
		}
		stds[x] = math.Sqrt(a / float64(n))
	}

	// 列参考值和列参考标准差
	var mu float64
	var sig float64
	for x := 0; x < probes; x++ {
		mu += means[x]
		sig += stds[x]
	}
	mu = mu / float64(probes)
	sig = sig / float64(probes)

	log.Printf("mean reference value: %f, std reference value: %f", mu, sig)
	// 修正 DN_adjusted[i] = (DN[i] - means[i]) *sig/stds[i]+mu
	for x := 0; x < originalImg.Cols(); x++ {
		for y := 0; y < originalImg.Rows(); y++ {
			dn := uint16(originalImg.GetShortAt(y, x))
			dn_adjusted := uint16((float64(dn)-means[x])*sig/stds[x] + mu)
			originalImg.SetShortAt(y, x, int16(dn_adjusted))
		}
	}
}

// 相邻列均衡算法（张兵，2006）- 实测结果不如上面的矩匹配算法
func DoMomentMatching2006(originalImg gocv.Mat) {
	probes := originalImg.Cols()
	log.Printf("do moment matching for %d probes", probes)
	// 第i个探元的像元均值和标准差
	means := make([]float64, probes)
	stds := make([]float64, probes)
	// 计算每个探元的均值和标准差
	for x := 0; x < originalImg.Cols(); x++ {
		var total int64
		n := originalImg.Rows()
		var dn uint16
		for y := 0; y < n; y++ {
			dn = uint16(originalImg.GetShortAt(y, x))
			total += int64(dn)
		}
		means[x] = float64(total) / float64(n)

		var a float64
		for y := 0; y < n; y++ {
			dn = uint16(originalImg.GetShortAt(y, x))
			a += math.Pow(float64(dn)-means[x], 2)
		}
		stds[x] = math.Sqrt(a / float64(n))
	}

	// 修正 DN_adjusted[i] = (DN[i] - means[i]) *sig/stds[i]+mu
	for x := 1; x < originalImg.Cols()-1; x++ {
		// 列参考值和列参考标准差
		mu := (means[x-1] + means[x+1]/2 + means[x]) / 2
		sig := (stds[x-1] + stds[x+1]/2 + stds[x]) / 2
		for y := 0; y < originalImg.Rows(); y++ {
			dn := uint16(originalImg.GetShortAt(y, x))
			dn_adjusted := uint16((float64(dn)-means[x])*sig/stds[x] + mu)
			originalImg.SetShortAt(y, x, int16(dn_adjusted))
		}
	}
}