#ifndef KNNCOLLE_HNSW_UTILS_HPP
#define KNNCOLLE_HNSW_UTILS_HPP

#include "hnswlib/hnswalg.h"

#include <functional>
#include <string>
#include <type_traits>
#include <filesystem>

namespace knncolle_hnsw {

/**
 * @cond
 */
template<typename Input_>
using I = std::remove_cv_t<std::remove_reference_t<Input_> >;
/**
 * @endcond
 */

/**
 * Define a global function to preserve `HnswData_` type information when saving a prebuilt Hnsw index in `knncolle::Prebuilt::save()`.
 * Users should define their own function here to handle an `HnswData_` type that is unknown to `knncolle::get_numeric_type()`. 
 * The action of setting/unsetting the global function is not thread-safe and should be done in a serial section. 
 *
 * The sole argument of the global function is the same `dir` provided to `knncolle::Prebuilt::save()`.
 * The global function is generally expected to write information about `HnswData_` to files inside `dir`.
 * It is recommended that the names of such files should not start with an upper-case letter to avoid conflicts with files generated by `save()`.
 * 
 * @tparam HnswData_ Type of data in the HNSW index, usually floating-point.
 *
 * @return Reference to a global function for saving information about `HnswData_`.
 * By default, no global function is defined.
 * If set, the global function will be called by the `knncolle::Prebuilt::save()` method for the HNSW `Prebuilt` subclass.
 */
template<class HnswData_>
std::function<void(const std::filesystem::path&)>& custom_save_for_hnsw_data() {
    static std::function<void(const std::filesystem::path&)> fun;
    return fun;
}

/**
 * Define a global function to save extra information about a distance metric of a prebuilt HNSW index in `knncolle::Prebuilt::save()`.
 * Users should define their own function here to handle a subclass of `hnswlib::SpaceInterface` that are unknown to `get_distance_name()`. 
 * The action of setting/unsetting the global function is not thread-safe and should be done in a serial section. 
 *
 * The first argument of the global function is the same `dir` provided to `knncolle::Prebuilt::save()`.
 * The second argument of the global function is a pointer to the distance metric used by the HNSW index.
 * The global function is generally expected to write information about the distance metric to files inside `dir`.
 * It is recommended that the names of such files should not start with an upper-case letter to avoid conflicts with files generated by `save()`.
 * 
 * The global function is only used when the index to be saved uses an unknown distance metric,
 * otherwise it is ignored.
 *
 * @tparam HnswData_ Type of data in the HNSW index, usually floating-point.
 *
 * @return Reference to a global function for saving information about the distance metric.
 * By default, no global function is defined.
 * If set, the global function will be called by the `knncolle::Prebuilt::save()` method for the HNSW `Prebuilt` subclass.
 */
template<typename HnswData_>
std::function<void(const std::filesystem::path&, const hnswlib::SpaceInterface<HnswData_>*)>& custom_save_for_hnsw_distance() {
    static std::function<void(const std::filesystem::path&, const hnswlib::SpaceInterface<HnswData_>*)> fun;
    return fun;
}

/**
 * Define a global function to save a custom distance normalization method for a prebuilt HNSW index.
 * Users should define their own function here to handle a `DistanceConfig::custom_normalize` function.
 * The action of setting/unsetting the global function is not thread-safe and should be done in a serial section. 
 *
 * The first argument of the global function is the same `dir` provided to `knncolle::Prebuilt::save()`.
 * The second argument of the global function is the `DistanceConfig::custom_normalize` method used to construct the HNSW index.
 * The global function is generally expected to write information about the normalization method to files inside `dir`.
 * It is recommended that the names of such files should not start with an upper-case letter to avoid conflicts with files generated by `save()`.
 *
 * The global function is only used when the index to be saved uses a custom distance normalization method,
 * otherwise it is ignored.
 * 
 * @tparam HnswData_ Type of data in the HNSW index, usually floating-point.
 *
 * @return Reference to a global function for saving information about the distance metric.
 * By default, no global function is defined.
 * If set, the global function will be called by the `knncolle::Prebuilt::save()` method for the HNSW `Prebuilt` subclass.
 */
template<typename Distance_>
std::function<void(const std::filesystem::path&, const std::function<Distance_(Distance_)>&)>& custom_save_for_hnsw_normalize() {
    static std::function<void(const std::filesystem::path&, const std::function<Distance_(Distance_)>&)> fun;
    return fun;
}

/**
 * Define a global function to create a HNSW distance metric that is not known to `get_distance_name()` when loading a HNSW index from disk.
 * Users are expected to provide their own function to regenerate any distance metric that was saved by `custom_save_for_hnsw_distance()`. 
 * Any modifications to this function are not thread-safe and should be done in a serial section. 
 *
 * The first argument to the global function is a directory path `dir`, same as that used in `knncolle::Prebuilt::save()`.
 * The second argument should be the number of dimensions in the dataset.
 * The global function should return a pointer to a `hnswlib::SpaceInterface` instance,
 * similar to the behavior of `DistanceConfig::create()`.
 *
 * The global function will only be called if the HNSW distance metric saved to `dir` is unknown to `get_distance_name()`,
 * otherwise it will ignored.
 * An error is raised if no global function is set when attempting to load an index with an unknown distance metric.
 * 
 * @tparam HnswData_ Type of data in the HNSW index, usually floating-point.
 *
 * @return Reference to a global function for creating a custom HNSW distance metric.
 * By default, no function is provided.
 * If set, the function will be called by the `knncolle::Prebuilt::save()` method for the HNSW `Prebuilt` subclass.
 */
template<typename HnswData_>
std::function<hnswlib::SpaceInterface<HnswData_>*(const std::filesystem::path&, std::size_t)>& custom_load_for_hnsw_distance() {
    static std::function<hnswlib::SpaceInterface<HnswData_>*(const std::filesystem::path&, std::size_t)> fun;
    return fun;
}

/**
 * Define a global function to create a custom distance normalization function when loading a HNSW index from disk.
 * Users are expected to provide their own function to regenerate any normalization function that was saved by `custom_save_for_hnsw_normalize()`. 
 * Any modifications to this function are not thread-safe and should be done in a serial section. 
 *
 * The sole argument to the global function is a directory path `dir`, same as that used in `knncolle::Prebuilt::save()`.
 * The global function should return a normalization function equivalent to the `DistanceConfig::custom_normalize()`.
 *
 * The global function is only used when the `DistanceConfig::normalize_method` saved to `dir` is equal to `DistanceNormalizeMethod::CUSTOM`.
 * otherwise it is ignored.
 * An error is raised if no global function is set when attempting to load an index with a custom normalization function.
 * 
 * @tparam Distance_ Floating point type for the distances.
 *
 * @return Reference to a global function for creating a custom distance metric.
 * By default, no function is provided.
 * If set, the function will be called by the `knncolle::Prebuilt::save()` method for the HNSW `Prebuilt` subclass.
 */
template<typename Distance_>
std::function<std::function<Distance_(Distance_)>(const std::filesystem::path&)>& custom_load_for_hnsw_normalize() {
    static std::function<std::function<Distance_(Distance_)>(const std::filesystem::path&)> fun;
    return fun;
}

}

#endif
