///|
pub fn kmeans(
  vectors : Array[Array[Double]],
  k : Int,
  max_iters : Int,
) -> Array[Array[Double]] raise VectorError {
  if k <= 0 {
    raise InvalidK
  }
  if vectors.length() == 0 {
    raise EmptyVector
  }
  if vectors.length() < k {
    raise ClusterError(
      "Cannot cluster " +
      vectors.length().to_string() +
      " points into " +
      k.to_string() +
      " clusters",
    )
  }

  let dim = vectors[0].length()
  for i = 1; i < vectors.length(); i = i + 1 {
    if vectors[i].length() != dim {
      raise DimensionMismatch("Vector dimension mismatch in K-Means clustering")
    }
  }

  // Initialize centroids with the first k vectors
  let centroids = Array::make(k, [])
  for i = 0; i < k; i = i + 1 {
    let copy = Array::make(dim, 0.0)
    for j = 0; j < dim; j = j + 1 {
      copy[j] = vectors[i][j]
    }
    centroids[i] = copy
  }

  let assignments = Array::make(vectors.length(), 0)
  for iter = 0; iter < max_iters; iter = iter + 1 {
    let mut changed = false
    for i = 0; i < vectors.length(); i = i + 1 {
      let mut min_dist = -1.0
      let mut best_centroid = 0
      for c = 0; c < k; c = c + 1 {
        let dist = euclidean_distance(vectors[i], centroids[c])
        if min_dist < 0.0 || dist < min_dist {
          min_dist = dist
          best_centroid = c
        }
      }
      if assignments[i] != best_centroid {
        assignments[i] = best_centroid
        changed = true
      }
    }

    if !changed && iter > 0 {
      break
    }

    let cluster_sizes = Array::make(k, 0)
    let new_centroids = Array::make(k, [])
    for c = 0; c < k; c = c + 1 {
      new_centroids[c] = Array::make(dim, 0.0)
    }

    for i = 0; i < vectors.length(); i = i + 1 {
      let c = assignments[i]
      cluster_sizes[c] = cluster_sizes[c] + 1
      for d = 0; d < dim; d = d + 1 {
        new_centroids[c][d] = new_centroids[c][d] + vectors[i][d]
      }
    }

    for c = 0; c < k; c = c + 1 {
      let size = cluster_sizes[c]
      if size > 0 {
        for d = 0; d < dim; d = d + 1 {
          centroids[c][d] = new_centroids[c][d] / size.to_double()
        }
      } else {
        let pick_idx = c % vectors.length()
        for d = 0; d < dim; d = d + 1 {
          centroids[c][d] = vectors[pick_idx][d]
        }
      }
    }
  }

  centroids
}