Kmknn.hpp

#ifndef KNNCOLLE_KMKNN_KMKNN_HPP
#define KNNCOLLE_KMKNN_KMKNN_HPP
 
#include "utils.hpp"
 
#include "knncolle/knncolle.hpp"
#include "kmeans/kmeans.hpp"
#include "sanisizer/sanisizer.hpp"
 
#include <algorithm>
#include <vector>
#include <memory>
#include <limits>
#include <cmath>
#include <cstddef>
#include <type_traits>
#include <string>
#include <filesystem>
 
namespace knncolle_kmknn {
 
inline static constexpr const char* kmknn_prebuilt_save_name = "knncolle_kmknn::Kmknn";
 
template<class KmeansFloat_>
std::function<void(const std::filesystem::path&)>& custom_save_for_kmknn_kmeansfloat() {
    static std::function<void(const std::filesystem::path&)> fun;
    return fun;
}
 
template<
    typename Index_,
    typename Data_,
    typename Distance_,
    typename KmeansIndex_ = Index_,
    typename KmeansData_ = Data_,
    typename KmeansCluster_ = Index_,
    typename KmeansFloat_ = Distance_,
    class KmeansMatrix_ = kmeans::SimpleMatrix<KmeansIndex_, KmeansData_>
>
struct KmknnOptions {
    double power = 0.5;
 
    std::shared_ptr<kmeans::Initialize<KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_> > initialize_algorithm;
 
    std::shared_ptr<kmeans::Refine<KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_> > refine_algorithm;
};
 
template<typename Index_, typename Data_, typename Distance_, class DistanceMetricData_, class KmeansFloat_, class DistanceMetricCenter_>
class KmknnPrebuilt;
 
template<typename Index_, typename Data_, typename Distance_, class DistanceMetricData_, class KmeansFloat_, class DistanceMetricCenter_>
class KmknnSearcher final : public knncolle::Searcher<Index_, Data_, Distance_> {
public:
    KmknnSearcher(const KmknnPrebuilt<Index_, Data_, Distance_, DistanceMetricData_, KmeansFloat_, DistanceMetricCenter_>& parent) : my_parent(parent) {
        my_center_order.reserve(my_parent.my_sizes.size());
        if constexpr(needs_conversion) {
            sanisizer::resize(my_query_conversion_buffer, my_parent.my_dim);
        }
    }
 
private:                
    const KmknnPrebuilt<Index_, Data_, Distance_, DistanceMetricData_, KmeansFloat_, DistanceMetricCenter_>& my_parent;
    knncolle::NeighborQueue<Index_, Distance_> my_nearest;
    std::vector<std::pair<Distance_, Index_> > my_all_neighbors;
    std::vector<std::pair<Distance_, Index_> > my_center_order;
 
    // Converting Data_ to KmeansFloat_ if we need to.
    static constexpr bool needs_conversion = !std::is_same<KmeansFloat_, Data_>::value;
    typename std::conditional<needs_conversion, std::vector<KmeansFloat_>, bool>::type my_query_conversion_buffer; 
 
    const KmeansFloat_* sanitize_query(const Data_* query) {
        if constexpr(needs_conversion) {
            auto conv_buffer = my_query_conversion_buffer.data();
            std::copy_n(query, my_parent.my_dim, conv_buffer);
            return conv_buffer;
        } else {
            return query;
        }
    }
 
    void finalize(std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) const {
        if (output_indices) {
            for (auto& s : *output_indices) {
                s = my_parent.my_observation_id[s];
            }
        }
        if (output_distances) {
            for (auto& d : *output_distances) {
                d = my_parent.my_metric_data->normalize(d);
            }
        }
    }
 
private:
    void search_nn(const Data_* query) {
        // Computing distances to all centers and sorting them.
        // The aim is to go through the nearest centers first, to try to get the shortest threshold (i.e., 'nearest.limit()') possible at the start;
        // this allows us to skip searches of the later clusters.
        {
            const auto query_san = sanitize_query(query);
            const auto ncenters = my_parent.my_sizes.size();
            my_center_order.clear();
            my_center_order.reserve(ncenters);
 
            for (I<decltype(ncenters)> c = 0; c < ncenters; ++c) {
                auto clust_ptr = my_parent.my_centers.data() + sanisizer::product_unsafe<std::size_t>(c, my_parent.my_dim);
                my_center_order.emplace_back(my_parent.my_metric_center->raw(my_parent.my_dim, query_san, clust_ptr), c);
            }
            std::sort(my_center_order.begin(), my_center_order.end());
        }
 
        // Computing the distance to each center, and deciding whether to proceed for each cluster.
        const auto& dist2centers = my_parent.my_dist_to_centroid;
        Distance_ threshold_raw = std::numeric_limits<Distance_>::infinity();
 
        for (const auto& curcent : my_center_order) {
            const Index_ center = curcent.second;
            Index_ firstsubj = my_parent.my_offsets[center], lastsubj = firstsubj + my_parent.my_sizes[center];
 
            if (!std::isinf(threshold_raw)) {
                const Distance_ threshold = my_parent.my_metric_center->normalize(threshold_raw);
                const Distance_ query2center = my_parent.my_metric_center->normalize(curcent.first);
                const Distance_ max_subj2center = dist2centers[lastsubj - 1];
 
                /* This exploits the triangle inequality to ignore points where:
                 *     threshold + subject-to-center < query-to-center 
                 * All points (if any) within this cluster with distances at or above 'lower_bd' are potentially countable.
                 *
                 * If the maximum distance between a subject and the center is less than 'lower_bd', there's no point proceeding,
                 * as we know that all other subjects will have smaller distances and are thus uncountable.
                 */
                const Distance_ lower_bd = query2center - threshold;
                if (max_subj2center < lower_bd) {
                    continue;
                }
                firstsubj = std::lower_bound(dist2centers.begin() + firstsubj, dist2centers.begin() + lastsubj, lower_bd) - dist2centers.begin();
 
                /* This exploits the reverse triangle inequality, to ignore points where:
                 *     threshold + query-to-center < subject-to-center
                 * All points (if any) within this cluster with distances at or below 'upper_bd' are potentially countable.
                 *
                 * If the maximum distance between a subject and the center is less than or equal to 'upper_bd', we can just skip the search.
                 * No subjects will a distance-to-center greater than 'upper_bd' so we know that we have to examine all subjects. 
                 *
                 * We could also skip this center altogther if the minimum subject-to-center distance is greater than 'upper_bd'.
                 * However, this seems too unlikely to warrant a special clause.
                 */
                const Distance_ upper_bd = query2center + threshold;
                if (max_subj2center > upper_bd) {
                    lastsubj = std::upper_bound(dist2centers.begin() + firstsubj, dist2centers.begin() + lastsubj, upper_bd) - dist2centers.begin();
                }
            }
 
            for (auto s = firstsubj; s < lastsubj; ++s) {
                const auto other_subj = my_parent.my_data.data() + sanisizer::product_unsafe<std::size_t>(s, my_parent.my_dim);
                auto dist2subj_raw = my_parent.my_metric_data->raw(my_parent.my_dim, query, other_subj);
                if (dist2subj_raw <= threshold_raw) {
                    my_nearest.add(s, dist2subj_raw);
                    if (my_nearest.is_full()) {
                        threshold_raw = my_nearest.limit(); // Shrinking the threshold, if an earlier NN has been found.
 
                        /* P.S. We could also consider increasing 'firstsubj' as 'threshold_raw' decreases. 
                         * The idea would be to exploit the triangle inequality to quickly skip over more points. 
                         * However, this is pointless because 'lower_bd' will never increase enough to skip subsequent observations.
                         * We wouldn't have been able to skip the observation that we just added,
                         * so there's no way we could skip observations with larger subject-to-center distances.
                         *
                         * P.P.S. We could also consider decreasing 'lastsubj' as 'threshold_raw' decreases.
                         * The idea would be to exploit the triangle inequality to terminate sooner. 
                         * However, this doesn't seem to provide a lot of benefit in practice. 
                         * In theory, we can only trim the search space if the query already lies in a center's hypersphere (as 'upper_bd' cannot decrease below 'query2center').
                         * Even then, 'upper_bd' is usually too large; testing indicates that a reduced 'upper_bd' only trims away a single observation at a time.
                         * There are also practical challenges as changes to 'lastsubj' within the loop might prevent out-of-order CPU execution;
                         * we need to do more memory accesses to 'dist2centers' to check if 'lastsubj' can be decreased;
                         * and we need to run an extra 'normalize()' to recompute 'upper_bd' inside the loop.
                         * All in all, I don't think it's worth it.
                         */
                    }
                }
            }
        }
    }
 
public:
    void search(Index_ i, Index_ k, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        my_nearest.reset(k + 1); // +1 is safe as k < num_obs.
        auto new_i = my_parent.my_new_location[i];
        auto iptr = my_parent.my_data.data() + sanisizer::product_unsafe<std::size_t>(new_i, my_parent.my_dim);
        search_nn(iptr);
        my_nearest.report(output_indices, output_distances, new_i);
        finalize(output_indices, output_distances);
    }
 
    void search(const Data_* query, Index_ k, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        if (k == 0) { // protect the NeighborQueue from k = 0.
            if (output_indices) {
                output_indices->clear();
            }
            if (output_distances) {
                output_distances->clear();
            }
        } else {
            my_nearest.reset(k);
            search_nn(query);
            my_nearest.report(output_indices, output_distances);
            finalize(output_indices, output_distances);
        }
    }
 
private:
    template<bool count_only_, typename Output_>
    void search_all(const Data_* query, Distance_ threshold, Output_& all_neighbors) {
        Distance_ threshold_raw = my_parent.my_metric_center->denormalize(threshold);
        const auto query_san = sanitize_query(query);
 
        // Computing distances to all centers. We don't sort them here because the threshold is constant so there's no point.
        const auto ncenters = my_parent.my_sizes.size();
        const auto& dist2centers = my_parent.my_dist_to_centroid;
 
        for (I<decltype(ncenters)> center = 0; center < ncenters; ++center) {
            auto center_ptr = my_parent.my_centers.data() + sanisizer::product_unsafe<std::size_t>(center, my_parent.my_dim);
            const Distance_ query2center = my_parent.my_metric_center->normalize(my_parent.my_metric_center->raw(my_parent.my_dim, query_san, center_ptr));
            Index_ firstsubj = my_parent.my_offsets[center], lastsubj = firstsubj + my_parent.my_sizes[center];
            const Distance_ max_subj2center = dist2centers[lastsubj - 1];
 
            // Same logic as in search_nn().
            const Distance_ lower_bd = query2center - threshold;
            if (max_subj2center < lower_bd) {
                continue;
            }
            firstsubj = std::lower_bound(dist2centers.begin() + firstsubj, dist2centers.begin() + lastsubj, lower_bd) - dist2centers.begin();
 
            // Same logic as in search_nn().
            const Distance_ upper_bd = query2center + threshold;
            if (max_subj2center > upper_bd) {
                lastsubj = std::upper_bound(dist2centers.begin() + firstsubj, dist2centers.begin() + lastsubj, upper_bd) - dist2centers.begin();
            }
 
            for (auto s = firstsubj; s < lastsubj; ++s) {
                const auto other_ptr = my_parent.my_data.data() + sanisizer::product_unsafe<std::size_t>(s, my_parent.my_dim);
                auto dist2cell_raw = my_parent.my_metric_data->raw(my_parent.my_dim, query, other_ptr);
                if (dist2cell_raw <= threshold_raw) {
                    if constexpr(count_only_) {
                        ++all_neighbors;
                    } else {
                        all_neighbors.emplace_back(dist2cell_raw, s);
                    }
                }
            }
        }
    }
 
public:
    bool can_search_all() const {
        return true;
    }
 
    Index_ search_all(Index_ i, Distance_ d, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        auto new_i = my_parent.my_new_location[i];
        auto iptr = my_parent.my_data.data() + sanisizer::product_unsafe<std::size_t>(new_i, my_parent.my_dim);
 
        if (!output_indices && !output_distances) {
            Index_ count = 0;
            search_all<true>(iptr, d, count);
            return knncolle::count_all_neighbors_without_self(count);
 
        } else {
            my_all_neighbors.clear();
            search_all<false>(iptr, d, my_all_neighbors);
            knncolle::report_all_neighbors(my_all_neighbors, output_indices, output_distances, new_i);
            finalize(output_indices, output_distances);
            return knncolle::count_all_neighbors_without_self(my_all_neighbors.size());
        }
    }
 
    Index_ search_all(const Data_* query, Distance_ d, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        if (!output_indices && !output_distances) {
            Index_ count = 0;
            search_all<true>(query, d, count);
            return count;
 
        } else {
            my_all_neighbors.clear();
            search_all<false>(query, d, my_all_neighbors);
            knncolle::report_all_neighbors(my_all_neighbors, output_indices, output_distances);
            finalize(output_indices, output_distances);
            return my_all_neighbors.size();
        }
    }
};
 
template<typename Index_, typename Data_, typename Distance_, class DistanceMetricData_, typename KmeansFloat_, class DistanceMetricCenter_>
class KmknnPrebuilt final : public knncolle::Prebuilt<Index_, Data_, Distance_> {
private:
    std::size_t my_dim;
    Index_ my_obs;
 
public:
    Index_ num_observations() const {
        return my_obs;
    }
 
    std::size_t num_dimensions() const {
        return my_dim;
    }
 
private:
    std::vector<Data_> my_data;
    std::shared_ptr<const DistanceMetricData_> my_metric_data;
    std::shared_ptr<const DistanceMetricCenter_> my_metric_center;
    
    std::vector<Index_> my_sizes;
    std::vector<Index_> my_offsets;
 
    std::vector<KmeansFloat_> my_centers;
 
    std::vector<Index_> my_observation_id, my_new_location;
    std::vector<Distance_> my_dist_to_centroid;
 
public:
    template<typename KmeansIndex_, typename KmeansData_, typename KmeansCluster_, class KmeansMatrix_>
    KmknnPrebuilt(
        std::size_t num_dim,
        Index_ num_obs,
        std::vector<Data_> data,
        std::shared_ptr<const DistanceMetricData_> metric_data,
        std::shared_ptr<const DistanceMetricCenter_> metric_center,
        const KmknnOptions<Index_, Data_, Distance_, KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_>& options
    ) :
        my_dim(num_dim),
        my_obs(num_obs),
        my_data(std::move(data)),
        my_metric_data(std::move(metric_data)),
        my_metric_center(std::move(metric_center))
    { 
        auto init = options.initialize_algorithm;
        if (init == nullptr) {
            init.reset(new kmeans::InitializeKmeanspp<KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_>);
        }
        auto refine = options.refine_algorithm;
        if (refine == nullptr) {
            refine.reset(new kmeans::RefineHartiganWong<KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_>);
        }
 
        KmeansCluster_ ncenters = sanisizer::from_float<KmeansCluster_>(std::ceil(std::pow(my_obs, options.power)));
        my_centers.resize(sanisizer::product<I<decltype(my_centers.size())> >(sanisizer::attest_gez(ncenters), my_dim));
 
        constexpr bool same_data = std::is_same<Data_, KmeansData_>::value;
        typename std::conditional<same_data, bool, std::vector<KmeansData_> >::type kmeans_data_buffer;
        const KmeansData_* data_ptr = NULL;
        if constexpr(same_data) {
            data_ptr = my_data.data();
        } else {
            kmeans_data_buffer.insert(kmeans_data_buffer.end(), my_data.begin(), my_data.end());
            data_ptr = kmeans_data_buffer.data();
        }
 
        kmeans::SimpleMatrix<KmeansIndex_, KmeansData_> mat(my_dim, sanisizer::cast<KmeansIndex_>(sanisizer::attest_gez(my_obs)), data_ptr);
        auto clusters = sanisizer::create<std::vector<KmeansCluster_> >(sanisizer::attest_gez(my_obs));
        auto output = kmeans::compute(mat, *init, *refine, ncenters, my_centers.data(), clusters.data());
 
        // Removing empty clusters, e.g., due to duplicate points.
        const auto survivors = kmeans::remove_unused_centers(my_dim, static_cast<KmeansIndex_>(my_obs), clusters.data(), ncenters, my_centers.data(), output.sizes);
        if (survivors < ncenters) {
            ncenters = survivors;
            my_centers.resize(sanisizer::product_unsafe<I<decltype(my_centers.size())> >(ncenters, my_dim));
            output.sizes.resize(ncenters);
        }
 
        if constexpr(std::is_same<Index_, KmeansIndex_>::value) {
            my_sizes.swap(output.sizes);
        } else {
            my_sizes.insert(my_sizes.end(), output.sizes.begin(), output.sizes.end());
        }
 
        sanisizer::resize(my_offsets, sanisizer::attest_gez(ncenters));
        for (KmeansCluster_ i = 1; i < ncenters; ++i) {
            my_offsets[i] = my_offsets[i - 1] + my_sizes[i - 1];
        }
 
        // Organize points correctly; firstly, sorting by distance from the assigned center.
        auto by_distance = sanisizer::create<std::vector<std::pair<Distance_, Index_> > >(sanisizer::attest_gez(my_obs));
        {
            static constexpr bool needs_conversion = !std::is_same<KmeansFloat_, Data_>::value;
            typename std::conditional<needs_conversion, std::vector<KmeansFloat_>, bool>::type conversion_buffer; 
            if constexpr(needs_conversion) {
                sanisizer::resize(conversion_buffer, my_dim);
            }
 
            auto sofar = my_offsets;
            for (Index_ o = 0; o < my_obs; ++o) {
                auto optr = my_data.data() + sanisizer::product_unsafe<std::size_t>(o, my_dim);
 
                const KmeansFloat_* observation = NULL;
                if constexpr(needs_conversion) {
                    std::copy_n(optr, my_dim, conversion_buffer.data());
                    observation = conversion_buffer.data();
                } else {
                    observation = optr;
                }
 
                auto clustid = clusters[o];
                auto cptr = my_centers.data() + sanisizer::product_unsafe<std::size_t>(clustid, my_dim);
 
                auto& counter = sofar[clustid];
                auto& current = by_distance[counter];
                current.first = my_metric_center->normalize(my_metric_center->raw(my_dim, observation, cptr));
                current.second = o;
 
                ++counter;
            }
 
            for (KmeansCluster_ c = 0; c < ncenters; ++c) {
                auto begin = by_distance.data() + my_offsets[c];
                std::sort(begin, begin + my_sizes[c]);
            }
        }
 
        // Permuting in-place to mirror the reordered distances, so that the search is more cache-friendly.
        {
            auto used = sanisizer::create<std::vector<unsigned char> >(sanisizer::attest_gez(my_obs));
            auto buffer = sanisizer::create<std::vector<Data_> >(my_dim);
            sanisizer::resize(my_observation_id, sanisizer::attest_gez(my_obs));
            sanisizer::resize(my_dist_to_centroid, sanisizer::attest_gez(my_obs));
            sanisizer::resize(my_new_location, sanisizer::attest_gez(my_obs));
 
            for (Index_ o = 0; o < my_obs; ++o) {
                if (used[o]) {
                    continue;
                }
 
                const auto& current = by_distance[o];
                my_observation_id[o] = current.second;
                my_dist_to_centroid[o] = current.first;
                my_new_location[current.second] = o;
                if (current.second == o) {
                    continue;
                }
 
                // We recursively perform a "thread" of replacements until we
                // are able to find the home of the originally replaced 'o'.
                auto optr = my_data.data() + sanisizer::product_unsafe<std::size_t>(o, my_dim);
                std::copy_n(optr, my_dim, buffer.data());
                Index_ replacement = current.second;
                do {
                    auto rptr = my_data.data() + sanisizer::product_unsafe<std::size_t>(replacement, my_dim);
                    std::copy_n(rptr, my_dim, optr);
                    used[replacement] = 1;
 
                    const auto& next = by_distance[replacement];
                    my_observation_id[replacement] = next.second;
                    my_dist_to_centroid[replacement] = next.first;
                    my_new_location[next.second] = replacement;
 
                    optr = rptr;
                    replacement = next.second;
                } while (replacement != o);
 
                std::copy(buffer.begin(), buffer.end(), optr);
            }
        }
    }
 
    friend class KmknnSearcher<Index_, Data_, Distance_, DistanceMetricData_, KmeansFloat_, DistanceMetricCenter_>;
 
public:
    std::unique_ptr<knncolle::Searcher<Index_, Data_, Distance_> > initialize() const {
        return initialize_known();
    }
 
    auto initialize_known() const {
        return std::make_unique<KmknnSearcher<Index_, Data_, Distance_, DistanceMetricData_, KmeansFloat_, DistanceMetricCenter_> >(*this);
    }
 
public:
    void save(const std::filesystem::path& dir) const {
        knncolle::quick_save(dir / "ALGORITHM", kmknn_prebuilt_save_name, std::strlen(kmknn_prebuilt_save_name));
        knncolle::quick_save(dir / "DATA", my_data.data(), my_data.size());
        knncolle::quick_save(dir / "NUM_OBS", &my_obs, 1);
        knncolle::quick_save(dir / "NUM_DIM", &my_dim, 1);
        const auto num_centers = my_sizes.size();
        knncolle::quick_save(dir / "NUM_CENTERS", &num_centers, 1);
 
        knncolle::quick_save(dir / "SIZES", my_sizes.data(), my_sizes.size());
        knncolle::quick_save(dir / "OFFSETS", my_offsets.data(), my_offsets.size());
        knncolle::quick_save(dir / "CENTERS", my_centers.data(), my_centers.size());
        knncolle::quick_save(dir / "OBSERVATION_ID", my_observation_id.data(), my_observation_id.size());
        knncolle::quick_save(dir / "NEW_LOCATION", my_new_location.data(), my_new_location.size());
        knncolle::quick_save(dir / "DIST_TO_CENTROID", my_dist_to_centroid.data(), my_dist_to_centroid.size());
 
        auto float_type = knncolle::get_numeric_type<KmeansFloat_>();
        knncolle::quick_save(dir / "FLOAT_TYPE", &float_type, 1);
        auto& kfcust = custom_save_for_kmknn_kmeansfloat<KmeansFloat_>(); 
        if (kfcust) {
            kfcust(dir); 
        }
 
        {
            const auto distdir = dir / "DISTANCE_DATA";
            std::filesystem::create_directory(distdir);
            my_metric_data->save(distdir);
        }
 
        {
            const auto distdir = dir / "DISTANCE_CENTER";
            std::filesystem::create_directory(distdir);
            my_metric_center->save(distdir);
        }
    }
 
    KmknnPrebuilt(const std::filesystem::path& dir) {
        knncolle::quick_load(dir / "NUM_OBS", &my_obs, 1);
        knncolle::quick_load(dir / "NUM_DIM", &my_dim, 1);
        auto num_centers = my_sizes.size();
        knncolle::quick_load(dir / "NUM_CENTERS", &num_centers, 1);
 
        my_data.resize(sanisizer::product<I<decltype(my_data.size())> >(sanisizer::attest_gez(my_obs), my_dim));
        knncolle::quick_load(dir / "DATA", my_data.data(), my_data.size());
 
        sanisizer::resize(my_sizes, sanisizer::attest_gez(num_centers));
        knncolle::quick_load(dir / "SIZES", my_sizes.data(), my_sizes.size());
        sanisizer::resize(my_offsets, sanisizer::attest_gez(num_centers));
        knncolle::quick_load(dir / "OFFSETS", my_offsets.data(), my_offsets.size());
        my_centers.resize(sanisizer::product<I<decltype(my_centers.size())> >(my_dim, sanisizer::attest_gez(num_centers)));
        knncolle::quick_load(dir / "CENTERS", my_centers.data(), my_centers.size());
 
        sanisizer::resize(my_observation_id, sanisizer::attest_gez(my_obs));
        knncolle::quick_load(dir / "OBSERVATION_ID", my_observation_id.data(), my_observation_id.size());
        sanisizer::resize(my_new_location, sanisizer::attest_gez(my_obs));
        knncolle::quick_load(dir / "NEW_LOCATION", my_new_location.data(), my_new_location.size());
        sanisizer::resize(my_dist_to_centroid, sanisizer::attest_gez(my_obs));
        knncolle::quick_load(dir / "DIST_TO_CENTROID", my_dist_to_centroid.data(), my_dist_to_centroid.size());
 
        {
            auto dptr = knncolle::load_distance_metric_raw<Data_, Distance_>(dir / "DISTANCE_DATA");
            auto xptr = dynamic_cast<DistanceMetricData_*>(dptr);
            if (xptr == NULL) {
                throw std::runtime_error("cannot cast the loaded distance metric to a DistanceMetricData_");
            }
            my_metric_data.reset(xptr);
        }
 
        {
            auto dptr = knncolle::load_distance_metric_raw<Data_, Distance_>(dir / "DISTANCE_CENTER");
            auto xptr = dynamic_cast<DistanceMetricCenter_*>(dptr);
            if (xptr == NULL) {
                throw std::runtime_error("cannot cast the loaded distance metric to a DistanceMetricCenter_");
            }
            my_metric_center.reset(xptr);
        }
    }
};
template<
    typename Index_,
    typename Data_,
    typename Distance_,
    class Matrix_ = knncolle::Matrix<Index_, Data_>,
    class DistanceMetricData_ = knncolle::DistanceMetric<Data_, Distance_>,
    typename KmeansIndex_ = Index_,
    typename KmeansData_ = Data_,
    typename KmeansCluster_ = Index_,
    typename KmeansFloat_ = Distance_,
    class KmeansMatrix_ = kmeans::SimpleMatrix<KmeansIndex_, KmeansData_>,
    class DistanceMetricCenter_ = knncolle::DistanceMetric<KmeansFloat_, Distance_>
>
class KmknnBuilder final : public knncolle::Builder<Index_, Data_, Distance_, Matrix_> {
public:
    typedef KmknnOptions<Index_, Data_, Distance_, KmeansIndex_, KmeansData_, KmeansCluster_, KmeansFloat_, KmeansMatrix_> Options;
 
private:
    std::shared_ptr<const DistanceMetricData_> my_metric_data;
    std::shared_ptr<const DistanceMetricCenter_> my_metric_center;
    Options my_options;
 
public:
    KmknnBuilder(
        std::shared_ptr<const DistanceMetricData_> metric_data,
        std::shared_ptr<const DistanceMetricCenter_> metric_center,
        Options options
    ) :
        my_metric_data(std::move(metric_data)),
        my_metric_center(std::move(metric_center)),
        my_options(std::move(options))
    {}
 
    KmknnBuilder(std::shared_ptr<const DistanceMetricData_> metric_data, std::shared_ptr<const DistanceMetricCenter_> metric_center) :
        KmknnBuilder(std::move(metric_data), std::move(metric_center), {}) {}
 
    // Don't provide an overload that accepts a raw metric pointer and the options,
    // as it's possible for the raw pointer to be constructed first, and then the options
    // is constructed but throws an error somewhere (e.g., in an IIFE), causing a memory leak.
    // as the raw pointer is never passed to the shared_ptr for management.
 
    Options& get_options() {
        return my_options;
    }
 
public:
    knncolle::Prebuilt<Index_, Data_, Distance_>* build_raw(const Matrix_& data) const {
        return build_known_raw(data);
    }
public:
    auto build_known_raw(const Matrix_& data) const {
        const auto ndim = data.num_dimensions();
        const auto nobs = data.num_observations();
 
        typedef std::vector<Data_> Store;
        Store store(sanisizer::product<typename Store::size_type>(ndim, nobs));
 
        auto work = data.new_known_extractor();
        for (I<decltype(nobs)> o = 0; o < nobs; ++o) {
            auto ptr = work->next();
            std::copy_n(ptr, ndim, store.data() + sanisizer::product_unsafe<std::size_t>(o, ndim)); 
        }
 
        return new KmknnPrebuilt<Index_, Data_, Distance_, DistanceMetricData_, KmeansFloat_, DistanceMetricCenter_>(
            ndim,
            nobs,
            std::move(store),
            my_metric_data,
            my_metric_center,
            my_options
        );
    }
 
    auto build_known_unique(const Matrix_& data) const {
        return std::unique_ptr<I<decltype(*build_known_raw(data))> >(build_known_raw(data));
    }
 
    auto build_known_shared(const Matrix_& data) const {
        return std::shared_ptr<I<decltype(*build_known_raw(data))> >(build_known_raw(data));
    }
};
 
}
 
#endif