// Requires the matching development runtime; published 0.0.5 uses legacy KMeans.
import aner.ml
import aner.parallel
import aner.tensor

let features = [[0.0], [2.0], [10.0], [12.0]]
let queries = [[1.0], [11.0]]
let serial = ml.fit_kmeans(features, 2, 100, 0.0, 42)
let execution = parallel.local()
let concurrent = ml.fit_kmeans(features: features, clusters: 2,
    max_iterations: 100, tolerance: 0.0, seed: 42,
    execution: execution, max_workspace_bytes: 64 * 1024 * 1024)

let serial_centers = serial.centers()
let concurrent_centers = concurrent.centers()
print(tensor.value(serial_centers, 0, 0) == tensor.value(concurrent_centers, 0, 0)
    && tensor.value(serial_centers, 1, 0) == tensor.value(concurrent_centers, 1, 0))
let serial_labels = serial.labels()
let concurrent_labels = concurrent.labels()
print(tensor.value(serial_labels, 0, 0) == tensor.value(concurrent_labels, 0, 0)
    && tensor.value(serial_labels, 1, 0) == tensor.value(concurrent_labels, 1, 0)
    && tensor.value(serial_labels, 2, 0) == tensor.value(concurrent_labels, 2, 0)
    && tensor.value(serial_labels, 3, 0) == tensor.value(concurrent_labels, 3, 0))
print(serial.inertia() == concurrent.inertia())
print(serial.iterations() == concurrent.iterations())
print(serial.converged() == concurrent.converged())
let serial_predictions = serial.predict(queries)
let concurrent_predictions = concurrent.predict(input: queries,
    execution: execution, max_workspace_bytes: 64 * 1024 * 1024)
print(tensor.value(serial_predictions, 0, 0) == tensor.value(concurrent_predictions, 0, 0)
    && tensor.value(serial_predictions, 1, 0) == tensor.value(concurrent_predictions, 1, 0))
print(concurrent.inertia()) // 4.0

// Exact accumulation preserves the small residual before dividing by three.
let balanced = ml.fit_kmeans([[1e150], [1.0], [-1e150]], 1, 3, 0.0, 42,
    execution: execution)
print(tensor.item(balanced.centers()) == 1.0 / 3.0)