package api

import (
	"image"
	"image/color"
	"math"
	"math/rand"
	"os"
)

// 定义一个结构来表示 RGB 颜色
type RGB struct {
	R, G, B int
}

// 计算两个 RGB 颜色的欧氏距离
func distance(c1, c2 RGB) float64 {
	return math.Sqrt(float64((c1.R-c2.R)*(c1.R-c2.R) + (c1.G-c2.G)*(c1.G-c2.G) + (c1.B-c2.B)*(c1.B-c2.B)))
}

// 对图像颜色进行 KMeans 聚类
func KMeans(colors []RGB, k int) ([]RGB, []int) {
	// 随机选择初始中心
	var centers []RGB
	for i := 0; i < k; i++ {
		centers = append(centers, colors[rand.Intn(len(colors))])
	}

	// 聚类算法
	labels := make([]int, len(colors))
	for iter := 0; iter < 10; iter++ { // 设置最大迭代次数为 10
		// 为每个颜色分配最近的聚类中心
		for i, color := range colors {
			minDist := distance(color, centers[0])
			labels[i] = 0
			for j := 1; j < k; j++ {
				dist := distance(color, centers[j])
				if dist < minDist {
					minDist = dist
					labels[i] = j
				}
			}
		}

		// 计算新的聚类中心
		newCenters := make([]RGB, k)
		counts := make([]int, k)
		for i, label := range labels {
			newCenters[label].R += colors[i].R
			newCenters[label].G += colors[i].G
			newCenters[label].B += colors[i].B
			counts[label]++
		}

		// 计算平均值作为新的中心
		for i := 0; i < k; i++ {
			if counts[i] > 0 {
				newCenters[i].R /= counts[i]
				newCenters[i].G /= counts[i]
				newCenters[i].B /= counts[i]
			}
		}

		// 如果中心没有变化，停止迭代
		changed := false
		for i := 0; i < k; i++ {
			if newCenters[i] != centers[i] {
				changed = true
				break
			}
		}
		if !changed {
			break
		}
		centers = newCenters
	}

	return centers, labels
}

// 加载图像并提取颜色
func extractColors(img image.Image) (colors []RGB) {
	bounds := img.Bounds()
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, g, b, _ := img.At(x, y).RGBA()
			colors = append(colors, RGB{int(r >> 8), int(g >> 8), int(b >> 8)})
		}
	}
	return colors
}

// 将聚类后的颜色应用到图像上
func recolorImage(img image.Image, centers []RGB, labels []int) image.Image {
	bounds := img.Bounds()
	newImage := image.NewRGBA(bounds)
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			label := labels[(y-bounds.Min.Y)*bounds.Max.X+(x-bounds.Min.X)]
			// 将 RGB 颜色转换为 RGBA 类型
			newColor := color.RGBA{
				R: uint8(centers[label].R),
				G: uint8(centers[label].G),
				B: uint8(centers[label].B),
				A: 255, // 不透明
			}
			// 设置新图像像素颜色
			newImage.Set(x, y, newColor)
		}
	}
	return newImage
}

func ImageToColors(str string) (colors []RGB, err error) {
	// 打开图像文件
	file, err := os.Open(str)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// 解码 JPEG 图像
	img, _, err := image.Decode(file)
	if err != nil {
		return nil, err
	}

	// 提取图像颜色
	centers, _ := KMeans(extractColors(img), 8)
	return centers, nil
}