naco-api/internal/clustering/kmeans_test.go

package clustering

import (
	"math"
	"math/rand"
	"testing"
)

func TestKMeans_Basic(t *testing.T) {
	items := []Item{
		{ID: "1", Text: "a", Embedding: []float64{1.0, 0.0}},
		{ID: "2", Text: "b", Embedding: []float64{0.9, 0.1}},
		{ID: "3", Text: "c", Embedding: []float64{-1.0, 0.0}},
		{ID: "4", Text: "d", Embedding: []float64{-0.9, -0.1}},
	}

	results, err := KMeans(items, 2, 50, 42)
	if err != nil {
		t.Fatalf("KMeans failed: %v", err)
	}

	if len(results) != 2 {
		t.Fatalf("expected 2 clusters, got %d", len(results))
	}

	// Check that items are assigned to clusters
	totalItems := 0
	for _, r := range results {
		totalItems += len(r.Items)
	}
	if totalItems != len(items) {
		t.Fatalf("expected %d items across clusters, got %d", len(items), totalItems)
	}

	// Verify centroids are normalized
	for _, r := range results {
		norm := 0.0
		for _, v := range r.Centroid {
			norm += v * v
		}
		norm = math.Sqrt(norm)
		if math.Abs(norm-1.0) > 1e-6 {
			t.Errorf("centroid not normalized: norm=%f", norm)
		}
	}
}

func TestKMeans_Empty(t *testing.T) {
	_, err := KMeans([]Item{}, 2, 50, 42)
	if err == nil {
		t.Fatal("expected error for empty items")
	}
}

func TestKMeans_KGreaterThanItems(t *testing.T) {
	items := []Item{
		{ID: "1", Text: "a", Embedding: []float64{1.0, 0.0}},
		{ID: "2", Text: "b", Embedding: []float64{0.0, 1.0}},
	}

	results, err := KMeans(items, 5, 50, 42)
	if err != nil {
		t.Fatalf("KMeans failed: %v", err)
	}

	// k should be capped to len(items)
	if len(results) != 2 {
		t.Fatalf("expected 2 clusters (capped), got %d", len(results))
	}
}

func TestKMeans_DeterministicWithSeed(t *testing.T) {
	items := []Item{
		{ID: "1", Text: "a", Embedding: []float64{1.0, 0.0}},
		{ID: "2", Text: "b", Embedding: []float64{0.9, 0.1}},
		{ID: "3", Text: "c", Embedding: []float64{-1.0, 0.0}},
		{ID: "4", Text: "d", Embedding: []float64{-0.9, -0.1}},
	}

	results1, _ := KMeans(items, 2, 50, 42)
	results2, _ := KMeans(items, 2, 50, 42)

	// Same seed should produce same assignments
	for i := range results1 {
		if len(results1[i].Items) != len(results2[i].Items) {
			t.Fatalf("cluster %d size differs: %d vs %d", i, len(results1[i].Items), len(results2[i].Items))
		}
	}
}

func TestCosineSimilarity(t *testing.T) {
	tests := []struct {
		name     string
		a, b     []float64
		expected float64
	}{
		{"identical", []float64{1, 0}, []float64{1, 0}, 1.0},
		{"orthogonal", []float64{1, 0}, []float64{0, 1}, 0.0},
		{"opposite", []float64{1, 0}, []float64{-1, 0}, -1.0},
		{"scaled", []float64{2, 0}, []float64{1, 0}, 2.0}, // not normalized
		{"mismatched len", []float64{1, 0}, []float64{1, 0, 0}, -1.0},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got := cosineSimilarity(tc.a, tc.b)
			if math.Abs(got-tc.expected) > 1e-6 {
				t.Errorf("cosineSimilarity(%v, %v) = %f, want %f", tc.a, tc.b, got, tc.expected)
			}
		})
	}
}

func TestNormalize(t *testing.T) {
	v := []float64{3.0, 4.0}
	normalized := normalize(v)

	// Check length is 1
	norm := 0.0
	for _, x := range normalized {
		norm += x * x
	}
	norm = math.Sqrt(norm)
	if math.Abs(norm-1.0) > 1e-6 {
		t.Errorf("normalized vector has norm %f, want 1.0", norm)
	}

	// Check zero vector handling
	zero := normalize([]float64{0.0, 0.0})
	if zero[0] != 0.0 || zero[1] != 0.0 {
		t.Errorf("zero vector normalization failed: %v", zero)
	}
}

func TestInitializeKMeansPlusPlus(t *testing.T) {
	vectors := [][]float64{
		{1.0, 0.0},
		{0.9, 0.1},
		{-1.0, 0.0},
		{-0.9, -0.1},
	}

	rng := rand.New(rand.NewSource(42))
	centroids := initializeKMeansPlusPlus(vectors, 2, rng)

	if len(centroids) != 2 {
		t.Fatalf("expected 2 centroids, got %d", len(centroids))
	}

	// Centroids should be from the input vectors (or averages)
	found1, found2 := false, false
	for _, c := range centroids {
		for _, v := range vectors {
			if equalVectors(c, v) {
				if !found1 {
					found1 = true
				} else {
					found2 = true
				}
			}
		}
	}
	if !found1 || !found2 {
		t.Errorf("centroids not from input vectors: %v", centroids)
	}
}

func equalVectors(a, b []float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if math.Abs(a[i]-b[i]) > 1e-6 {
			return false
		}
	}
	return true
}

func TestKMeans_HighDimensional(t *testing.T) {
	// Test with 1024-dim vectors (like qwen3-embedding)
	dim := 1024
	items := make([]Item, 10)
	for i := range items {
		emb := make([]float64, dim)
		if i < 5 {
			emb[0] = 1.0 // cluster 1
		} else {
			emb[dim-1] = 1.0 // cluster 2
		}
		items[i] = Item{
			ID:        string(rune('a' + i)),
			Text:      "test",
			Embedding: emb,
		}
	}

	results, err := KMeans(items, 2, 50, 42)
	if err != nil {
		t.Fatalf("KMeans failed: %v", err)
	}

	if len(results) != 2 {
		t.Fatalf("expected 2 clusters, got %d", len(results))
	}

	// Verify each cluster has ~5 items
	for _, r := range results {
		if len(r.Items) < 3 || len(r.Items) > 7 {
			t.Errorf("cluster size %d unexpected", len(r.Items))
		}
	}
}