#include "aner/random.hpp"
#include "aner/ml.hpp"
#include "aner/parallel_jobs.hpp"

#include <algorithm>
#include <atomic>
#include <exception>
#include <cmath>
#include <initializer_list>
#include <limits>
#include <numeric>
#include <optional>
#include <string>
#include <utility>
#include <vector>

namespace aner::ml {
namespace detail {
struct KNNState {
    std::size_t rows = 0;
    std::size_t cols = 0;
    std::size_t k = 0;
    bool distance_weights = false;
    std::vector<double> features;
    std::vector<double> classes;
    std::vector<std::size_t> row_classes;
};
struct KMeansState {
    std::size_t k = 0;
    std::size_t cols = 0;
    std::vector<double> centers;
    std::vector<double> labels;
    double inertia = 0;
    std::int64_t iterations = 0;
    bool converged = false;
};
struct Access {
    static const KNNState& get(const KNNClassifier& model) {
        if (!model.state_) throw Error("ml: KNNClassifier has not been fitted");
        return *model.state_;
    }
    static const KMeansState& get(const KMeansModel& model) {
        if (!model.state_) throw Error("ml: KMeansModel has not been fitted");
        return *model.state_;
    }
    static KNNClassifier make(std::shared_ptr<const KNNState> state) { return KNNClassifier(std::move(state)); }
    static KMeansModel make(std::shared_ptr<const KMeansState> state) { return KMeansModel(std::move(state)); }
};
} // namespace detail
namespace {
constexpr std::size_t max_elements = 1'000'000;
constexpr std::size_t max_work = 50'000'000;
constexpr double max_class_id = 9007199254740991.0;
struct TensorView {
    std::size_t rows;
    std::size_t cols;
    const std::vector<double>& values;
};
double finite(double value) {
    if (!std::isfinite(value)) throw Error("ml: nonfinite numeric input or unrepresentable intermediate result");
    return value;
}
std::size_t output_size(std::size_t rows, std::size_t cols) {
    if (rows == 0 || cols == 0 || rows > max_elements || cols > max_elements || rows > max_elements / cols)
        throw Error("ml: output exceeds the 1000000-element tensor limit");
    return rows * cols;
}
TensorView view(const nn::Tensor& tensor) {
    try {
        const TensorView result{tensor.rows(), tensor.cols(), tensor.values()};
        if (result.values.size() != output_size(result.rows, result.cols))
            throw Error("ml: tensor shape does not match its values");
        for (const auto value : result.values) finite(value);
        return result;
    } catch (const nn::Error& error) {
        throw Error(std::string("ml: invalid tensor: ") + error.what());
    }
}
void check_work(std::initializer_list<std::size_t> factors) {
    std::size_t product = 1;
    for (const auto factor : factors) {
        if (factor == 0 || factor > max_work / product)
            throw Error("ml: operation exceeds 50000000 distance-coordinate evaluations");
        product *= factor;
    }
}
nn::Tensor output(std::size_t rows, std::size_t cols, std::vector<double> values) {
    (void)output_size(rows, cols);
    try {
        return nn::tensor(static_cast<std::int64_t>(rows), static_cast<std::int64_t>(cols), std::move(values));
    } catch (const nn::Error& error) {
        throw Error(std::string("ml: invalid result: ") + error.what());
    }
}
void same_columns(std::size_t actual, std::size_t expected) {
    if (actual != expected) throw Error("ml: query columns must match fitted feature columns");
}
double distance(const double* a, const double* b, std::size_t cols) {
    double norm = 0;
    for (std::size_t col = 0; col < cols; ++col)
        norm = finite(std::hypot(norm, finite(a[col] - b[col])));
    return norm;
}
struct Neighbor { double distance; std::size_t index; };
bool neighbor_less(const Neighbor& a, const Neighbor& b) {
    return a.distance < b.distance || (a.distance == b.distance && a.index < b.index);
}
enum class KNNReadout { Prediction, Probabilities, Indices, Distances };
struct KNNLane {
    std::vector<Neighbor> neighbors;
    std::vector<std::size_t> votes;
    std::vector<double> minima, sums, corrections;
    std::size_t failed_row = std::numeric_limits<std::size_t>::max();
    std::exception_ptr failure;
};
struct KNNControl {
    std::atomic<std::size_t> first_failure;
    explicit KNNControl(std::size_t rows) : first_failure(rows) {}
};
class KNNWorkspace {
public:
    explicit KNNWorkspace(std::size_t limit) : limit_(limit) {
        if (!limit) throw Error("ml: KNN workspace budget must be positive");
    }
    void charge(std::size_t count, std::size_t width) {
        if (count > (limit_ - bytes_) / width)
            throw Error("ml: KNN output and lane workspace exceed max_workspace_bytes");
        bytes_ += count * width;
    }
private:
    std::size_t limit_, bytes_ = 0;
};
void compensated_add(double value, double& sum, double& correction) {
    const auto next = sum + value;
    correction += sum >= value ? (sum - next) + value : (value - next) + sum;
    sum = next;
}
bool knn_stopped(const KNNQueryOptions& options, const std::stop_token& executor_stop,
                 const KNNControl& control, std::size_t row) {
    return options.stop.stop_requested() || executor_stop.stop_requested() ||
           row > control.first_failure.load(std::memory_order_relaxed);
}
bool knn_distance(const double* a, const double* b, std::size_t cols,
                  const KNNQueryOptions& options, const std::stop_token& executor_stop,
                  const KNNControl& control, std::size_t row, double& norm) {
    norm = 0;
    for (std::size_t col = 0; col < cols; ++col) {
        if (col && (col & 255U) == 0 && knn_stopped(options, executor_stop, control, row)) return false;
        norm = finite(std::hypot(norm, finite(a[col] - b[col])));
    }
    return true;
}
std::size_t knn_vote(const detail::KNNState& fitted, KNNLane& lane,
                     double* probabilities) {
    const auto classes = fitted.classes.size();
    std::fill(lane.votes.begin(), lane.votes.end(), 0);
    if (!fitted.distance_weights || lane.neighbors[0].distance == 0) {
        std::size_t count = 0;
        for (std::size_t n = 0; n < fitted.k; ++n) {
            if (fitted.distance_weights && lane.neighbors[n].distance != 0) break;
            ++lane.votes[fitted.row_classes[lane.neighbors[n].index]];
            ++count;
        }
        const auto winner = static_cast<std::size_t>(std::max_element(lane.votes.begin(), lane.votes.end()) - lane.votes.begin());
        if (probabilities) for (std::size_t c = 0; c < classes; ++c)
            probabilities[c] = static_cast<double>(lane.votes[c]) / static_cast<double>(count);
        return winner;
    }
    std::fill(lane.minima.begin(), lane.minima.end(), std::numeric_limits<double>::infinity());
    std::fill(lane.sums.begin(), lane.sums.end(), 0);
    std::fill(lane.corrections.begin(), lane.corrections.end(), 0);
    // Sorted neighbors establish one canonical accumulation order. Class-local
    // scaling retains tiny class masses that individual global ratios could lose.
    for (std::size_t n = 0; n < fitted.k; ++n) {
        const auto& neighbor = lane.neighbors[n];
        const auto c = fitted.row_classes[neighbor.index];
        if (!std::isfinite(lane.minima[c])) lane.minima[c] = neighbor.distance;
        compensated_add(lane.minima[c] / neighbor.distance, lane.sums[c], lane.corrections[c]);
    }
    int nearest_exponent = 0;
    const auto nearest_mantissa = std::frexp(lane.neighbors[0].distance, &nearest_exponent);
    double total = 0, total_correction = 0;
    std::size_t winner = 0;
    for (std::size_t c = 0; c < classes; ++c) {
        double score = 0;
        if (std::isfinite(lane.minima[c])) {
            int class_exponent = 0;
            const auto class_mantissa = std::frexp(lane.minima[c], &class_exponent);
            const auto mantissa = (lane.sums[c] + lane.corrections[c]) * nearest_mantissa / class_mantissa;
            // Keep the unscaled mantissa for final normalization: rounding a
            // subnormal score before dividing can lose a representable result.
            lane.corrections[c] = mantissa;
            score = std::ldexp(mantissa, nearest_exponent - class_exponent);
        }
        lane.sums[c] = score;
        if (score > lane.sums[winner]) winner = c;
        compensated_add(score, total, total_correction);
    }
    total += total_correction;
    if (probabilities) for (std::size_t c = 0; c < classes; ++c) {
        if (!std::isfinite(lane.minima[c])) probabilities[c] = 0;
        else {
            int class_exponent = 0;
            (void)std::frexp(lane.minima[c], &class_exponent);
            probabilities[c] = std::ldexp(lane.corrections[c] / total, nearest_exponent - class_exponent);
        }
    }
    return winner;
}
nn::Tensor knn_query(const KNNClassifier& model, const nn::Tensor& queries,
                     const KNNQueryOptions& requested_options, KNNReadout readout) {
    // Native observers may rebind the caller's handles. Retain their immutable
    // owners and policy before exposing a callback; no numerical payload is cloned.
    const auto model_owner = model;
    const auto query_owner = queries;
    auto options = requested_options;
    std::optional<parallel::Context> execution_owner;
    if (options.execution) {
        execution_owner = *options.execution;
        options.execution = &*execution_owner;
    }
    if (options.stop.stop_requested()) throw Error("ml: KNN query cancelled");
    const auto& fitted = detail::Access::get(model_owner);
    const auto query = view(query_owner);
    same_columns(query.cols, fitted.cols);
    check_work({query.rows, fitted.rows, fitted.cols});
    std::size_t lanes = 1;
    parallel::Executor* executor = nullptr;
    if (options.execution) {
        const auto& context = *options.execution;
        if (context.nodes || context.device != "cpu" || context.dtype != "Float64" || !context.executor)
            throw Error("ml: KNN requires an explicit local CPU Float64 execution context with an executor");
        const auto workers = context.executor->workers();
        if (!workers || (context.workers && context.workers != workers))
            throw Error("ml: KNN execution context worker count does not match its executor");
        lanes = std::min(workers, query.rows);
        executor = context.executor.get();
    }
    const bool voting = readout == KNNReadout::Prediction || readout == KNNReadout::Probabilities;
    const bool weighted = voting && fitted.distance_weights;
    const auto columns = readout == KNNReadout::Prediction ? std::size_t{1} :
                         readout == KNNReadout::Probabilities ? fitted.classes.size() : fitted.k;
    const auto cells = output_size(query.rows, columns);
    // Admit the full operation once, then verify actual vector capacities after
    // allocation. Model/input caches and the existing executor are not workspace.
    auto charge = [&](KNNWorkspace& budget, std::size_t output_capacity,
                      std::size_t lane_capacity, const std::vector<KNNLane>* actual) {
        budget.charge(1, sizeof(KNNControl) + sizeof(std::vector<double>) + sizeof(std::vector<KNNLane>) +
                         sizeof(KNNClassifier) + sizeof(nn::Tensor) + sizeof(KNNQueryOptions) +
                         sizeof(std::optional<parallel::Context>));
        budget.charge(output_capacity, sizeof(double));
        budget.charge(lane_capacity, sizeof(KNNLane));
        for (std::size_t lane = 0; lane < lanes; ++lane) {
            const auto* scratch = actual ? &(*actual)[lane] : nullptr;
            budget.charge(scratch ? scratch->neighbors.capacity() : fitted.rows, sizeof(Neighbor));
            if (voting) budget.charge(scratch ? scratch->votes.capacity() : fitted.classes.size(), sizeof(std::size_t));
            if (weighted) {
                budget.charge(scratch ? scratch->minima.capacity() : fitted.classes.size(), sizeof(double));
                budget.charge(scratch ? scratch->sums.capacity() : fitted.classes.size(), sizeof(double));
                budget.charge(scratch ? scratch->corrections.capacity() : fitted.classes.size(), sizeof(double));
            }
        }
    };
    KNNWorkspace admission(options.max_workspace_bytes);
    charge(admission, cells, lanes, nullptr);
    std::vector<double> result;
    std::vector<KNNLane> scratch;
    try {
        result.resize(cells);
        scratch.resize(lanes);
        for (auto& lane : scratch) {
            lane.neighbors.resize(fitted.rows);
            if (voting) lane.votes.resize(fitted.classes.size());
            if (weighted) {
                lane.minima.resize(fitted.classes.size());
                lane.sums.resize(fitted.classes.size());
                lane.corrections.resize(fitted.classes.size());
            }
        }
    } catch (const std::bad_alloc&) { throw Error("ml: cannot allocate admitted KNN workspace"); }
    KNNWorkspace actual(options.max_workspace_bytes);
    charge(actual, result.capacity(), scratch.capacity(), &scratch);
    KNNControl control(query.rows);
    auto run_lane = [&](std::uint64_t id, std::stop_token executor_stop) {
        auto& lane = scratch[static_cast<std::size_t>(id)];
        const auto width = query.rows / lanes, remainder = query.rows % lanes;
        const auto first = static_cast<std::size_t>(id) * width + std::min(static_cast<std::size_t>(id), remainder);
        const auto last = first + width + (id < remainder ? 1 : 0);
        for (auto row = first; row < last; ++row) {
            if (knn_stopped(options, executor_stop, control, row)) break;
            try {
                bool complete = true;
                for (std::size_t training = 0; training < fitted.rows; ++training) {
                    if ((training & 63U) == 0 && knn_stopped(options, executor_stop, control, row)) { complete = false; break; }
                    double norm = 0;
                    if (!knn_distance(query.values.data() + row * fitted.cols,
                                      fitted.features.data() + training * fitted.cols, fitted.cols,
                                      options, executor_stop, control, row, norm)) { complete = false; break; }
                    lane.neighbors[training] = {norm, training};
                }
                if (!complete || knn_stopped(options, executor_stop, control, row)) break;
                if (fitted.k < lane.neighbors.size())
                    std::nth_element(lane.neighbors.begin(), lane.neighbors.begin() + static_cast<std::ptrdiff_t>(fitted.k),
                                     lane.neighbors.end(), neighbor_less);
                if (!voting || weighted)
                    std::sort(lane.neighbors.begin(), lane.neighbors.begin() + static_cast<std::ptrdiff_t>(fitted.k), neighbor_less);
                if (knn_stopped(options, executor_stop, control, row)) break;
                auto* destination = result.data() + row * columns;
                if (voting) {
                    const auto winner = knn_vote(fitted, lane, readout == KNNReadout::Probabilities ? destination : nullptr);
                    if (readout == KNNReadout::Prediction) destination[0] = fitted.classes[winner];
                } else for (std::size_t n = 0; n < fitted.k; ++n)
                    destination[n] = readout == KNNReadout::Indices ? static_cast<double>(lane.neighbors[n].index) : lane.neighbors[n].distance;
            } catch (...) {
                lane.failed_row = row;
                lane.failure = std::current_exception();
                auto earliest = control.first_failure.load(std::memory_order_relaxed);
                while (row < earliest && !control.first_failure.compare_exchange_weak(earliest, row, std::memory_order_relaxed)) {}
                break;
            }
        }
    };
    if (executor) {
        try {
            const auto summary = executor->run(lanes, run_lane, options.progress, options.stop);
            if (summary.cancelled || summary.completed_chunks != lanes)
                throw Error("ml: KNN query cancelled before all execution lanes completed");
        }
        catch (const parallel::Error& error) { throw Error(std::string("ml: KNN execution failed: ") + error.what()); }
    } else run_lane(0, {});
    if (options.stop.stop_requested()) throw Error("ml: KNN query cancelled");
    const auto failed = control.first_failure.load(std::memory_order_relaxed);
    if (failed < query.rows) for (const auto& lane : scratch) if (lane.failed_row == failed) {
        try { std::rethrow_exception(lane.failure); }
        catch (const std::exception& error) { throw Error("ml: KNN query row " + std::to_string(failed) + " failed: " + error.what()); }
        catch (...) { throw Error("ml: KNN query row " + std::to_string(failed) + " failed"); }
    }
    auto tensor = output(query.rows, columns, std::move(result));
    if (options.stop.stop_requested()) throw Error("ml: KNN query cancelled");
    return tensor;
}
std::size_t uniform_index(std::uint64_t& state, std::size_t count) {
    return static_cast<std::size_t>(rng::bounded(state, count));
}
std::pair<std::size_t, double> nearest_center(const double* point, const std::vector<double>& centers,
                                             std::size_t k, std::size_t cols) {
    std::size_t best = 0;
    auto best_distance = distance(point, centers.data(), cols);
    for (std::size_t cluster = 1; cluster < k; ++cluster) {
        const auto candidate = distance(point, centers.data() + cluster * cols, cols);
        if (candidate < best_distance) { best = cluster; best_distance = candidate; }
    }
    return {best, best_distance};
}
std::vector<std::size_t> assign(const TensorView& features, const std::vector<double>& centers, std::size_t k) {
    std::vector<std::size_t> result(features.rows);
    for (std::size_t row = 0; row < features.rows; ++row)
        result[row] = nearest_center(features.values.data() + row * features.cols, centers, k, features.cols).first;
    return result;
}
std::vector<double> update_centers(const TensorView& features, const std::vector<std::size_t>& assignment,
                                   const std::vector<double>& old_centers, std::size_t k) {
    std::vector<std::size_t> counts(k, 0);
    std::vector<double> magnitudes(old_centers.size(), 0);
    for (std::size_t row = 0; row < features.rows; ++row) {
        const auto cluster = assignment[row];
        ++counts[cluster];
        for (std::size_t col = 0; col < features.cols; ++col) {
            const auto index = cluster * features.cols + col;
            magnitudes[index] = std::max(magnitudes[index], std::abs(features.values[row * features.cols + col]));
        }
    }
    std::vector<double> sums(old_centers.size(), 0);
    std::vector<double> corrections(old_centers.size(), 0);
    for (std::size_t row = 0; row < features.rows; ++row)
        for (std::size_t col = 0; col < features.cols; ++col) {
            const auto index = assignment[row] * features.cols + col;
            if (magnitudes[index] == 0) continue;
            const auto term = features.values[row * features.cols + col] / magnitudes[index] - corrections[index];
            const auto next = sums[index] + term;
            corrections[index] = (next - sums[index]) - term;
            sums[index] = next;
        }
    auto result = old_centers;
    for (std::size_t cluster = 0; cluster < k; ++cluster) {
        if (counts[cluster] == 0) continue; // Deterministically retain empty clusters.
        for (std::size_t col = 0; col < features.cols; ++col) {
            const auto index = cluster * features.cols + col;
            result[index] = finite(magnitudes[index] * std::clamp(sums[index] / static_cast<double>(counts[cluster]), -1.0, 1.0));
        }
    }
    return result;
}
} // namespace

KNNClassifier::KNNClassifier(std::shared_ptr<const detail::KNNState> state) : state_(std::move(state)) {}
KMeansModel::KMeansModel(std::shared_ptr<const detail::KMeansState> state) : state_(std::move(state)) {}

KNNClassifier fit_knn(const nn::Tensor& features, const nn::Tensor& labels, std::int64_t k) {
    return fit_knn(features, labels, k, "uniform");
}
KNNClassifier fit_knn(const nn::Tensor& features, const nn::Tensor& labels, std::int64_t k,
                      const std::string& weights) {
    if (weights != "uniform" && weights != "distance")
        throw Error("ml.fit_knn: weights must be uniform or distance");
    const auto data = view(features);
    const auto target = view(labels);
    if (target.rows != data.rows || target.cols != 1)
        throw Error("ml.fit_knn: labels must have shape training rows x 1");
    if (k < 1 || static_cast<std::uint64_t>(k) > data.rows)
        throw Error("ml.fit_knn: k must be between one and the number of training rows");
    auto state = std::make_shared<detail::KNNState>();
    state->rows = data.rows;
    state->cols = data.cols;
    state->k = static_cast<std::size_t>(k);
    state->distance_weights = weights == "distance";
    state->features = data.values;
    state->classes = target.values;
    for (auto& label : state->classes) {
        if (label < 0 || label > max_class_id || std::trunc(label) != label)
            throw Error("ml.fit_knn: labels must be exact integers from zero through 9007199254740991");
        if (label == 0) label = 0; // Canonicalize negative zero.
    }
    std::sort(state->classes.begin(), state->classes.end());
    state->classes.erase(std::unique(state->classes.begin(), state->classes.end()), state->classes.end());
    state->row_classes.reserve(data.rows);
    for (const auto label : target.values)
        state->row_classes.push_back(static_cast<std::size_t>(std::lower_bound(state->classes.begin(), state->classes.end(), label) - state->classes.begin()));
    return detail::Access::make(std::move(state));
}
nn::Tensor predict(const KNNClassifier& model, const nn::Tensor& queries) {
    return predict(model, queries, KNNQueryOptions{});
}
nn::Tensor predict(const KNNClassifier& model, const nn::Tensor& queries, const KNNQueryOptions& options) {
    return knn_query(model, queries, options, KNNReadout::Prediction);
}
nn::Tensor predict_proba(const KNNClassifier& model, const nn::Tensor& queries) {
    return predict_proba(model, queries, KNNQueryOptions{});
}
nn::Tensor predict_proba(const KNNClassifier& model, const nn::Tensor& queries, const KNNQueryOptions& options) {
    return knn_query(model, queries, options, KNNReadout::Probabilities);
}
nn::Tensor classes(const KNNClassifier& model) {
    const auto& fitted = detail::Access::get(model);
    return output(fitted.classes.size(), 1, fitted.classes);
}
nn::Tensor neighbor_indices(const KNNClassifier& model, const nn::Tensor& queries) {
    return neighbor_indices(model, queries, KNNQueryOptions{});
}
nn::Tensor neighbor_indices(const KNNClassifier& model, const nn::Tensor& queries, const KNNQueryOptions& options) {
    return knn_query(model, queries, options, KNNReadout::Indices);
}
nn::Tensor neighbor_distances(const KNNClassifier& model, const nn::Tensor& queries) {
    return neighbor_distances(model, queries, KNNQueryOptions{});
}
nn::Tensor neighbor_distances(const KNNClassifier& model, const nn::Tensor& queries, const KNNQueryOptions& options) {
    return knn_query(model, queries, options, KNNReadout::Distances);
}

KMeansModel fit_kmeans(const nn::Tensor& features, std::int64_t k, std::int64_t max_iterations,
                       double tolerance, std::int64_t seed) {
    const auto data = view(features);
    if (k < 1 || static_cast<std::uint64_t>(k) > data.rows)
        throw Error("ml.fit_kmeans: k must be between one and the number of training rows");
    if (max_iterations < 1 || max_iterations > 1000)
        throw Error("ml.fit_kmeans: max_iterations must be between one and 1000");
    if (!std::isfinite(tolerance) || tolerance < 0)
        throw Error("ml.fit_kmeans: tolerance must be finite and nonnegative");
    const auto clusters = static_cast<std::size_t>(k);
    check_work({static_cast<std::size_t>(max_iterations) + 1, data.rows, clusters, data.cols});
    auto state = std::make_shared<detail::KMeansState>();
    state->k = clusters;
    state->cols = data.cols;
    state->centers.resize(output_size(clusters, data.cols));
    std::vector<std::size_t> indices(data.rows);
    std::iota(indices.begin(), indices.end(), 0);
    auto random_state = static_cast<std::uint64_t>(seed);
    for (std::size_t cluster = 0; cluster < clusters; ++cluster) {
        const auto selected = cluster + uniform_index(random_state, data.rows - cluster);
        std::swap(indices[cluster], indices[selected]);
        const auto start = data.values.begin() + static_cast<std::ptrdiff_t>(indices[cluster] * data.cols);
        std::copy_n(start, data.cols, state->centers.begin() + static_cast<std::ptrdiff_t>(cluster * data.cols));
    }
    std::vector<std::size_t> previous(data.rows, clusters);
    for (std::int64_t iteration = 0; iteration < max_iterations; ++iteration) {
        auto current = assign(data, state->centers, clusters);
        auto next = update_centers(data, current, state->centers, clusters);
        double movement = 0;
        for (std::size_t cluster = 0; cluster < clusters; ++cluster)
            movement = std::max(movement, distance(next.data() + cluster * data.cols,
                                                  state->centers.data() + cluster * data.cols, data.cols));
        state->centers = std::move(next);
        state->iterations = iteration + 1;
        if (current == previous || movement <= tolerance) {
            state->converged = true;
            break;
        }
        previous = std::move(current);
    }
    state->labels.resize(data.rows);
    for (std::size_t row = 0; row < data.rows; ++row) {
        const auto [cluster, dist] = nearest_center(data.values.data() + row * data.cols,
                                                  state->centers, clusters, data.cols);
        state->labels[row] = static_cast<double>(cluster);
        state->inertia = finite(state->inertia + finite(dist * dist));
    }
    return detail::Access::make(std::move(state));
}
nn::Tensor predict(const KMeansModel& model, const nn::Tensor& queries) {
    const auto& fitted = detail::Access::get(model);
    const auto query = view(queries);
    same_columns(query.cols, fitted.cols);
    check_work({query.rows, fitted.k, fitted.cols});
    std::vector<double> result(query.rows);
    for (std::size_t row = 0; row < query.rows; ++row)
        result[row] = static_cast<double>(nearest_center(query.values.data() + row * query.cols,
                                                        fitted.centers, fitted.k, fitted.cols).first);
    return output(query.rows, 1, std::move(result));
}
nn::Tensor centers(const KMeansModel& model) {
    const auto& fitted = detail::Access::get(model);
    return output(fitted.k, fitted.cols, fitted.centers);
}
nn::Tensor labels(const KMeansModel& model) {
    const auto& fitted = detail::Access::get(model);
    return output(fitted.labels.size(), 1, fitted.labels);
}
double inertia(const KMeansModel& model) { return detail::Access::get(model).inertia; }
std::int64_t iterations(const KMeansModel& model) { return detail::Access::get(model).iterations; }
bool converged(const KMeansModel& model) { return detail::Access::get(model).converged; }
} // namespace aner::ml