package fusion

import (
	"image"
	"os"
	"strings"

	"github.com/airbusgeo/godal"
	log "github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
	"starwiz.cn/sjy01/image-proc/pkg/config"
	"starwiz.cn/sjy01/image-proc/pkg/utils"
)

// Reference: https://www.kaggle.com/code/resolut/panchromatic-sharpening

type PansharpenMethod int

const (
	SimpleBrovey PansharpenMethod = iota
	Brovey
	SimpleMean
	ESRI
	GDALPansharpen
	IHS
	GIHS
	WeightR   = 1
	WeightG   = 0.75
	WeightB   = 0.55
	WeightNIR = 0.85
)

func Pansharpen(panTif, mssTif, fusTif string, method PansharpenMethod, weight float32) error {
	log.Infof("Pansharpening %s and %s", panTif, mssTif)
	pan, err := utils.ReadTifftoMat(panTif)
	if err != nil {
		return err
	}
	defer pan.Close()

	log.Infof("PAN row %d, col %d", pan.Rows(), pan.Cols())

	mss, err := utils.ReadTifftoMat(mssTif)
	if err != nil {
		return err
	}
	defer mss.Close()

	log.Infof("MSS bands %d row %d, col %d", mss.Channels(), mss.Rows(), mss.Cols())

	mssResized := gocv.NewMat()
	defer mssResized.Close()
	gocv.Resize(*mss, &mssResized, image.Point{X: pan.Cols(), Y: pan.Rows()}, 0, 0, gocv.InterpolationCubic)

	BGRI := gocv.Split(mssResized)

	var fus gocv.Mat
	switch method {
	case ESRI:
		fus = PansharpenESRI(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3])
	case IHS:
		fus = PansharpenIHS(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3], weight)
	case Brovey:
		fus = PansharpenBrovey(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3], weight)
	default:
		fus = PansharpenESRI(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3])
	}

	defer fus.Close()

	var mt gocv.MatType
	switch fus.Channels() {
	case 1:
		mt = gocv.MatTypeCV16UC1
	case 3:
		mt = gocv.MatTypeCV16UC3
	case 4:
		mt = gocv.MatTypeCV16UC4
	}
	fus.ConvertTo(&fus, mt)

	var colorInterps []godal.ColorInterp
	if config.GCONFIG.CoRegistration.FUSBandOrder == "RGB" ||
		config.GCONFIG.CoRegistration.FUSBandOrder == "RGBI" {
		colorInterps = []godal.ColorInterp{godal.CIRed, godal.CIGreen, godal.CIBlue, godal.CIUndefined}
	} else {
		colorInterps = []godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIUndefined}
	}

	err = utils.SaveBGRToGDALGTiff(fus, fus.Channels(), 0, 0, 1.3, colorInterps, fusTif)

	for i := 0; i < len(BGRI); i++ {
		BGRI[i].Close()
	}

	// use rpc of PAN directly
	rpb := strings.Replace(panTif, ".tiff", ".rpb", -1)
	if _, err := os.Stat(rpb); err == nil {
		fusRPB := strings.Replace(fusTif, ".tiff", ".rpb", -1)
		if input, err := os.ReadFile(rpb); err == nil {
			os.WriteFile(fusRPB, input, 0644)
		}
	}

	return err
}

func PansharpenSimpleBrovey(pan, B, G, R, I gocv.Mat) error {
	return nil
}

func PansharpenBrovey(pan, B, G, R, I gocv.Mat, w float32) gocv.Mat {
	pan.ConvertTo(&pan, gocv.MatTypeCV32F)
	B.ConvertTo(&B, gocv.MatTypeCV32F)
	G.ConvertTo(&G, gocv.MatTypeCV32F)
	R.ConvertTo(&R, gocv.MatTypeCV32F)
	I.ConvertTo(&I, gocv.MatTypeCV32F)

	// DNF = (PAN - w*I)/(R*w + G*w + B*w)
	wI := gocv.NewMat()
	I.CopyTo(&wI)
	wI.MultiplyFloat(-1 * w)
	defer wI.Close()
	wR := gocv.NewMat()
	R.CopyTo(&wR)
	wR.MultiplyFloat(w)
	defer wR.Close()
	wG := gocv.NewMat()
	G.CopyTo(&wG)
	wG.MultiplyFloat(w)
	defer wG.Close()
	wB := gocv.NewMat()
	B.CopyTo(&wB)
	wB.MultiplyFloat(w)
	defer wB.Close()

	DNF := gocv.NewMat()
	defer DNF.Close()
	gocv.Add(pan, wI, &DNF)
	gocv.Add(wR, wG, &wR)
	gocv.Add(wR, wB, &wR)
	gocv.Divide(DNF, wR, &DNF)

	gocv.Multiply(R, DNF, &R)
	gocv.Multiply(G, DNF, &G)
	gocv.Multiply(B, DNF, &B)
	gocv.Multiply(I, DNF, &I)

	mss := gocv.NewMat()
	mt := gocv.MatTypeCV16UC3
	switch config.GCONFIG.CoRegistration.FUSBandOrder {
	case "RGB":
		gocv.Merge([]gocv.Mat{R, G, B}, &mss)
	case "BGR":
		gocv.Merge([]gocv.Mat{B, G, R}, &mss)
	case "RGBI":
		gocv.Merge([]gocv.Mat{R, G, B, I}, &mss)
		mt = gocv.MatTypeCV16UC4
	case "BGRI":
		gocv.Merge([]gocv.Mat{B, G, R, I}, &mss)
		mt = gocv.MatTypeCV16UC4
	default:
		gocv.Merge([]gocv.Mat{R, G, B}, &mss)
	}

	mss.ConvertTo(&mss, mt)

	return mss
}

func PansharpenSimpleMean(pan, B, G, R, I gocv.Mat) error { return nil }

func PansharpenESRI(pan, B, G, R, I gocv.Mat) gocv.Mat {
	bgrn_mean := gocv.NewMatWithSize(pan.Rows(), pan.Cols(), gocv.MatTypeCV32F)
	defer bgrn_mean.Close()
	ADJ := gocv.NewMat()
	defer ADJ.Close()

	pan.ConvertTo(&pan, gocv.MatTypeCV32F)
	B.ConvertTo(&B, gocv.MatTypeCV32F)
	G.ConvertTo(&G, gocv.MatTypeCV32F)
	R.ConvertTo(&R, gocv.MatTypeCV32F)
	I.ConvertTo(&I, gocv.MatTypeCV32F)

	for y := 0; y < pan.Rows(); y++ {
		for x := 0; x < pan.Cols(); x++ {
			v := (R.GetFloatAt(y, x) + G.GetFloatAt(y, x) + B.GetFloatAt(y, x) + I.GetFloatAt(y, x)) / 4
			bgrn_mean.SetFloatAt(y, x, v)
		}
	}

	bgrn_mean.MultiplyFloat(-1)
	gocv.Add(pan, bgrn_mean, &ADJ)
	gocv.Add(B, ADJ, &B)
	gocv.Add(G, ADJ, &G)
	gocv.Add(R, ADJ, &R)
	gocv.Add(I, ADJ, &I)

	mss := gocv.NewMat()
	mt := gocv.MatTypeCV16UC3
	switch config.GCONFIG.CoRegistration.FUSBandOrder {
	case "RGB":
		gocv.Merge([]gocv.Mat{R, G, B}, &mss)
	case "BGR":
		gocv.Merge([]gocv.Mat{B, G, R}, &mss)
	case "RGBI":
		gocv.Merge([]gocv.Mat{R, G, B, I}, &mss)
		mt = gocv.MatTypeCV16UC4
	case "BGRI":
		gocv.Merge([]gocv.Mat{B, G, R, I}, &mss)
		mt = gocv.MatTypeCV16UC4
	default:
		gocv.Merge([]gocv.Mat{R, G, B}, &mss)
	}

	mss.ConvertTo(&mss, mt)

	return mss
}

func PansharpenIHS(pan, B, G, R, I gocv.Mat, w float32) gocv.Mat {
	mss := gocv.NewMat()
	gocv.Merge([]gocv.Mat{B, G, R}, &mss)
	mssHLS := gocv.NewMat()
	defer mssHLS.Close()

	mss.ConvertTo(&mss, gocv.MatTypeCV32FC3)
	gocv.CvtColor(mss, &mssHLS, gocv.ColorBGRToHLS)
	hlsChannels := gocv.Split(mssHLS)

	// 将 PAN 图像转换为与 HLS 亮度分量相同的范围
	panFloat := gocv.NewMat()
	defer panFloat.Close()
	pan.ConvertTo(&panFloat, gocv.MatTypeCV32F)
	gocv.Normalize(panFloat, &panFloat, 0, 65535, gocv.NormMinMax)

	hlsFloat := gocv.NewMat()
	defer hlsFloat.Close()
	hlsChannels[1].ConvertTo(&hlsFloat, gocv.MatTypeCV32F)
	gocv.Normalize(hlsFloat, &hlsFloat, 0, 65535, gocv.NormMinMax)

	// 替换亮度分量
	I.ConvertTo(&I, gocv.MatTypeCV32F)
	I.MultiplyFloat(-1 * w)
	gocv.Add(panFloat, I, &hlsChannels[1])

	// 合并通道
	gocv.Merge(hlsChannels, &mssHLS)

	// 将图像从 HLS 转换回 BGR
	fusedImage := gocv.NewMat()
	gocv.CvtColor(mssHLS, &fusedImage, gocv.ColorHLSToBGR)

	channels := gocv.Split(fusedImage)
	defer channels[0].Close() // B
	defer channels[1].Close() // G
	defer channels[2].Close() // R

	if config.GCONFIG.CoRegistration.FUSBandOrder == "RGBI" || config.GCONFIG.CoRegistration.FUSBandOrder == "RGB" {
		gocv.Merge([]gocv.Mat{channels[2], channels[1], channels[0]}, &fusedImage)
	} else {
		gocv.Merge([]gocv.Mat{channels[0], channels[1], channels[2]}, &fusedImage)
	}

	fusedImage.ConvertTo(&fusedImage, gocv.MatTypeCV16UC3)

	return fusedImage
}