package cloudcover

import (
	"image"
	"image/color"

	"math"
	"math/rand"
	"os"
	"time"

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

const (
	k                   = 2   // 聚类数目，2表示云和非云
	maxIters            = 100 // 最大迭代次数
	tolerance           = 1.0 // 聚类中心变化容忍度
	brightnessThreshold = 220 // 亮度阈值，值越大要求云更亮
)

type Pixel struct {
	r, g, b float64
}

func CloudPercentByKMeans(imagePath string) float64 {
	file, err := os.Open(imagePath)
	if err != nil {
		log.Errorf("无法打开图像: %v", err)
		return 0.0
	}
	defer file.Close()

	img, _, err := image.Decode(file)
	if err != nil {
		log.Errorf("无法解码图像: %v", err)
		return 0.0
	}

	return computeCloudCoverByKMeans(img)
}

func CloudPercentByKMeans2(imgMat gocv.Mat) float64 {
	img, err := imgMat.ToImage()
	if err != nil {
		log.Errorf("无法将Mat转换为Image: %v", err)
		return 0.0
	}
	return computeCloudCoverByKMeans(img)
}

func computeCloudCoverByKMeans(img image.Image) float64 {
	log.Info("cloud cover computing...")
	bounds := img.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y
	pixels := make([]Pixel, 0, width*height)

	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			r, g, b, _ := img.At(x, y).RGBA()
			pixels = append(pixels, Pixel{
				r: float64(r >> 8),
				g: float64(g >> 8),
				b: float64(b >> 8),
			})
		}
	}

	rand.Seed(time.Now().UnixNano())
	centroids := initializeCentroids(pixels)
	assignments := make([]int, len(pixels))

	for iter := 0; iter < maxIters; iter++ {
		for i, pixel := range pixels {
			assignments[i] = nearestCentroid(pixel, centroids)
		}

		newCentroids := updateCentroids(pixels, assignments)
		if hasConverged(centroids, newCentroids) {
			log.Debugf("Convergence after the %d iteration", iter+1)
			break
		}
		centroids = newCentroids
	}

	cloudPixels := 0
	totalPixels := width * height

	// 新增：检查哪个聚类的平均亮度更高，确保白色表示云
	cloudCluster := 0
	if centroids[1].r+centroids[1].g+centroids[1].b > centroids[0].r+centroids[0].g+centroids[0].b {
		cloudCluster = 1
	}

	outputImg := image.NewGray(bounds)
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			i := y*width + x
			pixel := pixels[i]
			// 计算亮度和颜色比值
			brightness := (pixel.r + pixel.g + pixel.b) / 3
			blueRedRatio := pixel.b / pixel.r
			greenRedRatio := pixel.g / pixel.r

			// 判断是否为云的条件：亮度大于阈值并且蓝色/红色比率较高
			if (assignments[i] == cloudCluster && brightness > brightnessThreshold && blueRedRatio > 0.5) ||
				(assignments[i] == cloudCluster && greenRedRatio > 0.8) {
				outputImg.SetGray(x, y, color.Gray{Y: 255}) // 云区域
				cloudPixels++
			} else {
				outputImg.SetGray(x, y, color.Gray{Y: 0}) // 非云区域
			}
		}
	}

	cloudRatio := float64(cloudPixels) / float64(totalPixels)
	return cloudRatio
}

func initializeCentroids(pixels []Pixel) []Pixel {
	centroids := make([]Pixel, k)
	for i := range centroids {
		centroids[i] = pixels[rand.Intn(len(pixels))]
	}
	return centroids
}

func nearestCentroid(pixel Pixel, centroids []Pixel) int {
	minDist := math.Inf(1)
	bestIndex := 0
	for i, c := range centroids {
		dist := math.Sqrt(math.Pow(pixel.r-c.r, 2) + math.Pow(pixel.g-c.g, 2) + math.Pow(pixel.b-c.b, 2))
		if dist < minDist {
			minDist = dist
			bestIndex = i
		}
	}
	return bestIndex
}

func updateCentroids(pixels []Pixel, assignments []int) []Pixel {
	sums := make([]Pixel, k)
	counts := make([]int, k)

	for i, assignment := range assignments {
		sums[assignment].r += pixels[i].r
		sums[assignment].g += pixels[i].g
		sums[assignment].b += pixels[i].b
		counts[assignment]++
	}

	newCentroids := make([]Pixel, k)
	for i := range newCentroids {
		if counts[i] > 0 {
			newCentroids[i].r = sums[i].r / float64(counts[i])
			newCentroids[i].g = sums[i].g / float64(counts[i])
			newCentroids[i].b = sums[i].b / float64(counts[i])
		} else {
			newCentroids[i] = sums[i]
		}
	}

	return newCentroids
}

func hasConverged(oldCentroids, newCentroids []Pixel) bool {
	for i := range oldCentroids {
		if math.Abs(oldCentroids[i].r-newCentroids[i].r) > tolerance ||
			math.Abs(oldCentroids[i].g-newCentroids[i].g) > tolerance ||
			math.Abs(oldCentroids[i].b-newCentroids[i].b) > tolerance {
			return false
		}
	}
	return true
}