package clustering

import (
	"errors"
	"math"
	"math/rand"
)

func initializeKMeansPlusPlus(
	vectors [][]float64,
	k int,
	rng *rand.Rand,
) [][]float64 {
	centroids := make([][]float64, 0, k)

	first := rng.Intn(len(vectors))
	centroids = append(
		centroids,
		append([]float64(nil), vectors[first]...),
	)

	distances := make([]float64, len(vectors))

	for len(centroids) < k {
		var total float64

		for i, vector := range vectors {
			bestSimilarity := -math.MaxFloat64

			for _, centroid := range centroids {
				similarity := cosineSimilarity(vector, centroid)
				if similarity > bestSimilarity {
					bestSimilarity = similarity
				}
			}

			// cosine distance
			distance := 1 - bestSimilarity

			if distance < 0 {
				distance = 0
			}

			// k-means++ uses squared distance.
			distances[i] = distance * distance
			total += distances[i]
		}

		if total == 0 {
			break
		}

		target := rng.Float64() * total

		var cumulative float64
		selected := len(vectors) - 1

		for i, distance := range distances {
			cumulative += distance

			if cumulative >= target {
				selected = i
				break
			}
		}

		centroids = append(
			centroids,
			append([]float64(nil), vectors[selected]...),
		)
	}

	// This can theoretically happen if all vectors are identical.
	for len(centroids) < k {
		index := rng.Intn(len(vectors))

		centroids = append(
			centroids,
			append([]float64(nil), vectors[index]...),
		)
	}

	return centroids
}

func KMeans(items []Item, k int, maxIterations int, seed int64) ([]Result, error) {
	if len(items) == 0 {
		return nil, errors.New("no items")
	}

	if k <= 0 {
		return nil, errors.New("k must be positive")
	}

	if k > len(items) {
		k = len(items)
	}

	if maxIterations <= 0 {
		maxIterations = 50
	}

	dimension := len(items[0].Embedding)
	if dimension == 0 {
		return nil, errors.New("empty embedding")
	}

	// Normalize embeddings once.
	normalized := make([][]float64, len(items))
	for i, item := range items {
		if len(item.Embedding) != dimension {
			return nil, errors.New("embedding dimensions do not match")
		}

		normalized[i] = normalize(item.Embedding)
	}

	rng := rand.New(rand.NewSource(seed))

	centroids := initializeKMeansPlusPlus(normalized, k, rng)

	assignments := make([]int, len(items))

	for iteration := 0; iteration < maxIterations; iteration++ {
		changed := false

		// Assignment step.
		for i, vector := range normalized {
			bestCluster := 0
			bestSimilarity := -math.MaxFloat64

			for cluster, centroid := range centroids {
				similarity := cosineSimilarity(vector, centroid)

				if similarity > bestSimilarity {
					bestSimilarity = similarity
					bestCluster = cluster
				}
			}

			if iteration == 0 || assignments[i] != bestCluster {
				assignments[i] = bestCluster
				changed = true
			}
		}

		// Stop if nothing changed.
		if !changed && iteration > 0 {
			break
		}

		// Update step.
		newCentroids := make([][]float64, k)
		counts := make([]int, k)

		for cluster := 0; cluster < k; cluster++ {
			newCentroids[cluster] = make([]float64, dimension)
		}

		for i, vector := range normalized {
			cluster := assignments[i]
			counts[cluster]++

			for d, value := range vector {
				newCentroids[cluster][d] += value
			}
		}

		for cluster := 0; cluster < k; cluster++ {
			if counts[cluster] == 0 {
				// Empty cluster: reinitialize with a random item.
				index := rng.Intn(len(normalized))
				newCentroids[cluster] = append(
					[]float64(nil),
					normalized[index]...,
				)
				continue
			}

			newCentroids[cluster] = normalize(newCentroids[cluster])
		}

		centroids = newCentroids
	}

	results := make([]Result, k)

	for cluster := 0; cluster < k; cluster++ {
		results[cluster] = Result{
			Cluster:  cluster,
			Centroid: centroids[cluster],
		}
	}

	for i, item := range items {
		cluster := assignments[i]

		item.Embedding = normalized[i]

		results[cluster].Items = append(
			results[cluster].Items,
			item,
		)
	}

	// Find the item closest to each centroid.
	for cluster := range results {
		if len(results[cluster].Items) == 0 {
			continue
		}

		bestSimilarity := -math.MaxFloat64

		for _, item := range results[cluster].Items {
			similarity := cosineSimilarity(
				item.Embedding,
				results[cluster].Centroid,
			)

			if similarity > bestSimilarity {
				bestSimilarity = similarity
				results[cluster].Representative = item
			}
		}
	}

	return results, nil
}