package producer

import (
	"gocv.io/x/gocv"
)

const (
	NoEnhancement                = "NoEnhancement"
	StretchToMinimumMaximum      = "StretchToMinimumMaximum"
	StretchToCumulativeCutMinMax = "StretchToCumulativeCutMinMax"
)

type ContrastEnhancement struct {
	minimumValue         int
	maximumValue         int
	lookupTableOffset    int
	lookupTable          []int
	mRasterDataTypeRange int
	enhancementType      string
	//!< Range is [ min + cumulativeCutLower() * (max - min), min + cumulativeCutUpper() * (max - min) ]
	cumulativeCutLower float64
	cumulativeCutUpper float64
}

func NewContrastEnhancement(minimumValue, maximumValue int) *ContrastEnhancement {
	ce := &ContrastEnhancement{
		enhancementType:    StretchToMinimumMaximum,
		cumulativeCutLower: 0,
		cumulativeCutUpper: 1,
	}
	ce.minimumValue = minimumValue + int(ce.cumulativeCutLower)*(maximumValue-minimumValue)
	ce.maximumValue = minimumValue + int(ce.cumulativeCutUpper)*(maximumValue-minimumValue)
	ce.mRasterDataTypeRange = ce.maximumValue - ce.minimumValue
	ce.lookupTableOffset = -1 * ce.minimumValue
	ce.lookupTable = make([]int, ce.mRasterDataTypeRange+1)
	ce.generateLookupTable()

	return ce
}

func (ce ContrastEnhancement) enhance(value int) int {
	if ce.enhancementType != NoEnhancement {
		shiftedValue := value + ce.lookupTableOffset
		if shiftedValue >= 0 && shiftedValue < ce.mRasterDataTypeRange+1 {
			return ce.lookupTable[shiftedValue]
		}
		return 0
	}

	return ce.linearMinMaxEnhancement(value)
}

func (ce *ContrastEnhancement) generateLookupTable() {
	for i := range ce.lookupTable {
		ce.lookupTable[i] = ce.linearMinMaxEnhancement(i - ce.lookupTableOffset)
	}
}

func (ce ContrastEnhancement) linearMinMaxEnhancement(value int) int {
	stretchedValue := (float64(value-ce.minimumValue) / float64(ce.mRasterDataTypeRange)) * 255.0
	if stretchedValue < 0 {
		return 0
	} else if stretchedValue > 255 {
		return 255
	}
	return int(stretchedValue)
}

func cumulativeCountCutEnhancement(img gocv.Mat, cumulativeCutLower, cumulativeCutUpper float64) gocv.Mat {
	dataTypeMax := 65536
	// 计算直方图
	hist := make([]int, dataTypeMax)
	for y := 0; y < img.Rows(); y++ {
		for x := 0; x < img.Cols(); x++ {
			hist[uint16(img.GetShortAt(y, x))]++
		}
	}

	// 计算累积直方图
	cdf := make([]int, dataTypeMax)
	cdf[0] = hist[0]
	for i := 1; i < dataTypeMax; i++ {
		cdf[i] = cdf[i-1] + hist[i]
	}

	// 计算低高阈值
	totalPixels := img.Rows() * img.Cols()
	lowCutoff := int(float64(totalPixels) * cumulativeCutLower)
	highCutoff := int(float64(totalPixels) * cumulativeCutUpper)

	var lowerBound, upperBound int
	for i, v := range cdf {
		if v > lowCutoff {
			lowerBound = i
			break
		}
	}
	for i := dataTypeMax - 1; i >= 0; i-- {
		if cdf[i] < highCutoff {
			upperBound = i
			break
		}
	}

	// 重新映射灰度值
	enhancedImg := gocv.NewMatWithSize(img.Rows(), img.Cols(), gocv.MatTypeCV8U)
	for y := 0; y < img.Rows(); y++ {
		for x := 0; x < img.Cols(); x++ {
			oldVal := uint16(img.GetShortAt(y, x))
			newVal := (int(oldVal) - lowerBound) * 255 / (upperBound - lowerBound)
			if newVal > 255 {
				newVal = 255
			}
			if newVal < 0 {
				newVal = 0
			}

			enhancedImg.SetUCharAt(y, x, uint8(newVal))
		}
	}

	return enhancedImg

}