mcpnet.cpp

#include "mcpnet.h"

#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>

#include <mcp/correlation/AdaptivePartitioningMI.hpp>
#include <mcp/correlation/BSplineMI.hpp>
#include <mcp/correlation/Pearson.hpp>
#include <mcp/filter/dpi.hpp>
#include <mcp/filter/mcp.hpp>
#include <mcp/filter/threshold.hpp>
#include <mcp/transform/clr.hpp>
#include <mcp/transform/combine.hpp>
#include <mcp/transform/stouffer.hpp>
#include <splash/transform/rank.hpp>
#include <splash/transform/zscore.hpp>

#include <utility>

// -----------------------------------------------------------------------------
// AdaptivePartitionRankMIKernel2
// -----------------------------------------------------------------------------

namespace {

template <typename OT>
inline OT adaptive_mi_kernel_impl(rust::Slice<const ::std::size_t> x_ranks,
                                  rust::Slice<const ::std::size_t> y_ranks) {
    using IT = ::std::size_t;

    const ::std::size_t count = x_ranks.size();
    if (count == 0 || count != y_ranks.size()) {
        return static_cast<OT>(0);
    }

    // Build cumulative histograms for the two rank vectors as a 2-row matrix.
    // The MI kernel expects (count + 1) entries per row and addresses rows
    // by a byte stride.
    const ::std::size_t hist_cols = count + 1;
    ::std::vector<IT> hist(2 * hist_cols, 0);

    using HistogramKernelType =
        ::mcp::kernel::IntegralCumulativeHistogramKernel<IT, 0>;
    HistogramKernelType hist_kernel;
    hist_kernel(x_ranks.data(), count, hist.data());
    hist_kernel(y_ranks.data(), count, hist.data() + hist_cols);

    using MIKernelType =
        ::mcp::kernel::AdaptivePartitionRankMIKernel2<IT, OT, 0>;
    MIKernelType mi_op(hist.data(), hist_cols * sizeof(IT));

    return mi_op(0, 1, x_ranks.data(), y_ranks.data(), count);
}

} // namespace

float adaptive_mi_kernel_f32(rust::Slice<const ::std::size_t> x_ranks,
                             rust::Slice<const ::std::size_t> y_ranks) {
    return adaptive_mi_kernel_impl<float>(x_ranks, y_ranks);
}

double adaptive_mi_kernel_f64(rust::Slice<const ::std::size_t> x_ranks,
                              rust::Slice<const ::std::size_t> y_ranks) {
    return adaptive_mi_kernel_impl<double>(x_ranks, y_ranks);
}

// -----------------------------------------------------------------------------
// Rank (descending, firstRank = 0, output = size_t -- matches mi.cpp)
// -----------------------------------------------------------------------------

namespace {

template <typename IT>
inline void rank_impl(rust::Slice<const IT>      in_vec,
                      rust::Slice<::std::size_t> out_vec) {
    if (in_vec.size() == 0) return;
    assert(out_vec.size() >= in_vec.size() &&
           "rank: output buffer too small");

    using RankKernelType = ::splash::kernel::Rank<IT, ::std::size_t, 0, false>;
    RankKernelType ranker;
    ranker(in_vec.data(), in_vec.size(), out_vec.data());
}

} // namespace

void rank_f32(rust::Slice<const float>  in_vec,
              rust::Slice<::std::size_t> out_vec) {
    rank_impl<float>(in_vec, out_vec);
}

void rank_f64(rust::Slice<const double> in_vec,
              rust::Slice<::std::size_t> out_vec) {
    rank_impl<double>(in_vec, out_vec);
}

// -----------------------------------------------------------------------------
// IntegralCumulativeHistogramKernel
// -----------------------------------------------------------------------------

void integral_cumulative_histogram(rust::Slice<const ::std::size_t> in_vec,
                                   rust::Slice<::std::size_t>       out_vec) {
    if (in_vec.size() == 0) return;
    assert(out_vec.size() >= in_vec.size() + 1 &&
           "integral_cumulative_histogram: output buffer must be at least "
           "in_vec.size() + 1 long");

    using IT = ::std::size_t;
    using HistogramKernelType =
        ::mcp::kernel::IntegralCumulativeHistogramKernel<IT, 0>;
    HistogramKernelType hist_kernel;
    hist_kernel(in_vec.data(), in_vec.size(), out_vec.data());
}

// -----------------------------------------------------------------------------
// BSplineWeightsKernel
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline void bspline_weights_impl(rust::Slice<const T> in_vec,
                                 ::std::int32_t       num_bins,
                                 ::std::int32_t       spline_order,
                                 ::std::int32_t       num_samples,
                                 rust::Slice<T>       out_vec) {
    assert(static_cast<::std::int32_t>(in_vec.size()) >= num_samples &&
           "bspline_weights: in_vec must hold at least num_samples elements");
    const ::std::size_t required =
        static_cast<::std::size_t>(num_bins) *
            static_cast<::std::size_t>(num_samples) + 1u;
    assert(out_vec.size() >= required &&
           "bspline_weights: out_vec must be at least num_bins*num_samples+1");

    using KernelType = ::mcp::kernel::BSplineWeightsKernel<T, T>;
    KernelType kernel(num_bins, spline_order, num_samples);
    kernel(in_vec.data(), in_vec.size(), out_vec.data());
}

} // namespace

void bspline_weights_f32(rust::Slice<const float> in_vec,
                         ::std::int32_t num_bins,
                         ::std::int32_t spline_order,
                         ::std::int32_t num_samples,
                         rust::Slice<float> out_vec) {
    bspline_weights_impl<float>(in_vec, num_bins, spline_order, num_samples,
                                out_vec);
}

void bspline_weights_f64(rust::Slice<const double> in_vec,
                         ::std::int32_t num_bins,
                         ::std::int32_t spline_order,
                         ::std::int32_t num_samples,
                         rust::Slice<double> out_vec) {
    bspline_weights_impl<double>(in_vec, num_bins, spline_order, num_samples,
                                 out_vec);
}

// -----------------------------------------------------------------------------
// BSplineMIKernel
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline T bspline_mi_kernel_impl(rust::Slice<const T> first,
                                rust::Slice<const T> second,
                                ::std::int32_t       num_bins,
                                ::std::int32_t       num_samples) {
    const ::std::size_t required =
        static_cast<::std::size_t>(num_bins) *
            static_cast<::std::size_t>(num_samples) + 1u;
    assert(first.size()  >= required &&
           "bspline_mi_kernel: first must be at least num_bins*num_samples+1");
    assert(second.size() >= required &&
           "bspline_mi_kernel: second must be at least num_bins*num_samples+1");

    using KernelType = ::mcp::kernel::BSplineMIKernel<T, T>;
    KernelType kernel(num_bins, num_samples);
    return kernel(first.data(), second.data(), required);
}

} // namespace

float bspline_mi_kernel_f32(rust::Slice<const float> first,
                            rust::Slice<const float> second,
                            ::std::int32_t num_bins,
                            ::std::int32_t num_samples) {
    return bspline_mi_kernel_impl<float>(first, second, num_bins, num_samples);
}

double bspline_mi_kernel_f64(rust::Slice<const double> first,
                             rust::Slice<const double> second,
                             ::std::int32_t num_bins,
                             ::std::int32_t num_samples) {
    return bspline_mi_kernel_impl<double>(first, second, num_bins, num_samples);
}

// -----------------------------------------------------------------------------
// StandardScore
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline void standard_score_impl(rust::Slice<const T> in_vec,
                                rust::Slice<T>       out_vec) {
    if (in_vec.size() == 0) return;
    assert(out_vec.size() >= in_vec.size() &&
           "standard_score: output buffer too small");

    // Use population variance (SampleStats=false) so that the result is
    // compatible with mcp::correlation::PearsonKernel, which divides the
    // dot product by N -- matches the convention in MCPNet's pearson.cpp.
    using KernelType = ::splash::kernel::StandardScore<T, T, false>;
    KernelType kernel;
    kernel(in_vec.data(), in_vec.size(), out_vec.data());
}

} // namespace

void standard_score_f32(rust::Slice<const float>  in_vec,
                        rust::Slice<float>        out_vec) {
    standard_score_impl<float>(in_vec, out_vec);
}

void standard_score_f64(rust::Slice<const double> in_vec,
                        rust::Slice<double>       out_vec) {
    standard_score_impl<double>(in_vec, out_vec);
}

// -----------------------------------------------------------------------------
// PearsonKernel
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline T pearson_kernel_impl(rust::Slice<const T> first,
                             rust::Slice<const T> second) {
    if (first.size() == 0 || first.size() != second.size()) {
        return static_cast<T>(0);
    }

    using KernelType = ::mcp::correlation::PearsonKernel<T, T>;
    KernelType kernel;
    return kernel(first.data(), second.data(), first.size());
}

} // namespace

float pearson_kernel_f32(rust::Slice<const float> first,
                         rust::Slice<const float> second) {
    return pearson_kernel_impl<float>(first, second);
}

double pearson_kernel_f64(rust::Slice<const double> first,
                          rust::Slice<const double> second) {
    return pearson_kernel_impl<double>(first, second);
}

// -----------------------------------------------------------------------------
// combine.hpp -- element-wise binary kernels
// -----------------------------------------------------------------------------

namespace {

template <typename Kernel, typename T>
inline void apply_binary(Kernel const & k,
                         rust::Slice<const T> a,
                         rust::Slice<const T> b,
                         rust::Slice<T>       out) {
    const ::std::size_t count =
        a.size() < b.size() ? a.size() : b.size();
    if (count == 0) return;
    assert(out.size() >= count &&
           "combine kernel: out buffer too small");
    k(a.data(), b.data(), count, out.data());
}

template <typename Kernel, typename T>
inline void apply_unary(Kernel const & k,
                        rust::Slice<const T> in,
                        rust::Slice<T>       out) {
    if (in.size() == 0) return;
    assert(out.size() >= in.size() &&
           "combine kernel: out buffer too small");
    k(in.data(), in.size(), out.data());
}

} // namespace

#define DEFINE_BINARY_KERNEL(NAME, KERNEL_TYPE)                              \
    void NAME ## _f32(rust::Slice<const float>  a,                            \
                      rust::Slice<const float>  b,                            \
                      rust::Slice<float>        out) {                        \
        ::mcp::kernel::KERNEL_TYPE<float, float> kernel;                      \
        apply_binary(kernel, a, b, out);                                      \
    }                                                                         \
    void NAME ## _f64(rust::Slice<const double> a,                            \
                      rust::Slice<const double> b,                            \
                      rust::Slice<double>       out) {                        \
        ::mcp::kernel::KERNEL_TYPE<double, double> kernel;                    \
        apply_binary(kernel, a, b, out);                                      \
    }

DEFINE_BINARY_KERNEL(add_kernel,      add_kernel)
DEFINE_BINARY_KERNEL(sub_kernel,      sub_kernel)
DEFINE_BINARY_KERNEL(multiply_kernel, multiply_kernel)
DEFINE_BINARY_KERNEL(ratio_kernel,    ratio_kernel)

#undef DEFINE_BINARY_KERNEL

// max_kernel/min_kernel are single-typed (IT only).
#define DEFINE_BINARY_KERNEL_SINGLE_T(NAME, KERNEL_TYPE)                      \
    void NAME ## _f32(rust::Slice<const float>  a,                            \
                      rust::Slice<const float>  b,                            \
                      rust::Slice<float>        out) {                        \
        ::mcp::kernel::KERNEL_TYPE<float> kernel;                             \
        apply_binary(kernel, a, b, out);                                      \
    }                                                                         \
    void NAME ## _f64(rust::Slice<const double> a,                            \
                      rust::Slice<const double> b,                            \
                      rust::Slice<double>       out) {                        \
        ::mcp::kernel::KERNEL_TYPE<double> kernel;                            \
        apply_binary(kernel, a, b, out);                                      \
    }

DEFINE_BINARY_KERNEL_SINGLE_T(max_kernel, max_kernel)
DEFINE_BINARY_KERNEL_SINGLE_T(min_kernel, min_kernel)

#undef DEFINE_BINARY_KERNEL_SINGLE_T

void madd_kernel_f32(rust::Slice<const float> a, rust::Slice<const float> b,
                     double coeff, rust::Slice<float> out) {
    ::mcp::kernel::madd_kernel<float, float> kernel(coeff);
    apply_binary(kernel, a, b, out);
}
void madd_kernel_f64(rust::Slice<const double> a, rust::Slice<const double> b,
                     double coeff, rust::Slice<double> out) {
    ::mcp::kernel::madd_kernel<double, double> kernel(coeff);
    apply_binary(kernel, a, b, out);
}

void weighted_mean_kernel_f32(rust::Slice<const float> a,
                              rust::Slice<const float> b,
                              double alpha, double beta,
                              rust::Slice<float> out) {
    ::mcp::kernel::weighted_mean_kernel<float, float> kernel(alpha, beta);
    apply_binary(kernel, a, b, out);
}
void weighted_mean_kernel_f64(rust::Slice<const double> a,
                              rust::Slice<const double> b,
                              double alpha, double beta,
                              rust::Slice<double> out) {
    ::mcp::kernel::weighted_mean_kernel<double, double> kernel(alpha, beta);
    apply_binary(kernel, a, b, out);
}

void scale_kernel_f32(rust::Slice<const float> in, double coeff,
                      rust::Slice<float> out) {
    ::mcp::kernel::scale_kernel<float, float> kernel(coeff);
    apply_unary(kernel, in, out);
}
void scale_kernel_f64(rust::Slice<const double> in, double coeff,
                      rust::Slice<double> out) {
    ::mcp::kernel::scale_kernel<double, double> kernel(coeff);
    apply_unary(kernel, in, out);
}

// -----------------------------------------------------------------------------
// stouffer.hpp / clr.hpp -- helpers
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline ::std::vector<::std::pair<T, T>> zip_pairs(
    rust::Slice<const T> means,
    rust::Slice<const T> stds) {
    const ::std::size_t n = means.size() < stds.size() ? means.size() : stds.size();
    ::std::vector<::std::pair<T, T>> pairs(n);
    for (::std::size_t i = 0; i < n; ++i) {
        pairs[i] = ::std::make_pair(means[i], stds[i]);
    }
    return pairs;
}

template <typename Kernel, typename T>
inline T scalar_pair_pair(Kernel const & k,
                          T value,
                          T mean1, T std1,
                          T mean2, T std2) {
    ::std::pair<T, T> a1{mean1, std1};
    ::std::pair<T, T> a2{mean2, std2};
    return k(value, a1, a2);
}

template <typename Kernel, typename T>
inline void apply_pair_vector(Kernel const & k,
                              rust::Slice<const T> in_vec,
                              T aux1_mean, T aux1_std,
                              rust::Slice<const T> aux2_means,
                              rust::Slice<const T> aux2_stds,
                              rust::Slice<T> out) {
    const ::std::size_t count = in_vec.size();
    if (count == 0) return;
    assert(aux2_means.size() >= count &&
           "pair-vector kernel: aux2_means too small");
    assert(aux2_stds.size()  >= count &&
           "pair-vector kernel: aux2_stds too small");
    assert(out.size() >= count &&
           "pair-vector kernel: out buffer too small");

    ::std::pair<T, T> aux1{aux1_mean, aux1_std};
    auto aux2_pairs = zip_pairs<T>(aux2_means, aux2_stds);
    k(in_vec.data(), &aux1, aux2_pairs.data(), count, out.data());
}

template <typename Kernel, typename T>
inline void apply_zscored_binary(Kernel const & k,
                                 rust::Slice<const T> zi,
                                 rust::Slice<const T> zj,
                                 rust::Slice<T> out) {
    const ::std::size_t count =
        zi.size() < zj.size() ? zi.size() : zj.size();
    if (count == 0) return;
    assert(out.size() >= count &&
           "zscored kernel: out buffer too small");
    k(zi.data(), zj.data(), count, out.data());
}

} // namespace

// ---- Stouffer ---------------------------------------------------------------

float stouffer_kernel_f32(float value,
                          float m1, float s1, float m2, float s2) {
    ::mcp::kernel::stouffer_kernel<float, ::std::pair<float, float>, float> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}
double stouffer_kernel_f64(double value,
                           double m1, double s1, double m2, double s2) {
    ::mcp::kernel::stouffer_kernel<double, ::std::pair<double, double>, double> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}

void stouffer_vector_kernel_f32(rust::Slice<const float> in_vec,
                                float aux1_mean, float aux1_std,
                                rust::Slice<const float> aux2_means,
                                rust::Slice<const float> aux2_stds,
                                rust::Slice<float> out) {
    ::mcp::kernel::stouffer_vector_kernel<float, ::std::pair<float, float>, float> k;
    apply_pair_vector(k, in_vec, aux1_mean, aux1_std,
                      aux2_means, aux2_stds, out);
}
void stouffer_vector_kernel_f64(rust::Slice<const double> in_vec,
                                double aux1_mean, double aux1_std,
                                rust::Slice<const double> aux2_means,
                                rust::Slice<const double> aux2_stds,
                                rust::Slice<double> out) {
    ::mcp::kernel::stouffer_vector_kernel<double, ::std::pair<double, double>, double> k;
    apply_pair_vector(k, in_vec, aux1_mean, aux1_std,
                      aux2_means, aux2_stds, out);
}

void zscored_stouffer_kernel_f32(rust::Slice<const float> zi,
                                 rust::Slice<const float> zj,
                                 rust::Slice<float> out) {
    ::mcp::kernel::zscored_stouffer_kernel<float, float, float> k;
    apply_zscored_binary(k, zi, zj, out);
}
void zscored_stouffer_kernel_f64(rust::Slice<const double> zi,
                                 rust::Slice<const double> zj,
                                 rust::Slice<double> out) {
    ::mcp::kernel::zscored_stouffer_kernel<double, double, double> k;
    apply_zscored_binary(k, zi, zj, out);
}

float weighted_stouffer_kernel_f32(float value,
                                   float m1, float s1, float m2, float s2) {
    ::mcp::kernel::weighted_stouffer_kernel<float, ::std::pair<float, float>, float> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}
double weighted_stouffer_kernel_f64(double value,
                                    double m1, double s1, double m2, double s2) {
    ::mcp::kernel::weighted_stouffer_kernel<double, ::std::pair<double, double>, double> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}

// ---- CLR --------------------------------------------------------------------

float clr_kernel_f32(float value,
                     float m1, float s1, float m2, float s2) {
    ::mcp::kernel::clr_kernel<float, ::std::pair<float, float>, float> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}
double clr_kernel_f64(double value,
                     double m1, double s1, double m2, double s2) {
    ::mcp::kernel::clr_kernel<double, ::std::pair<double, double>, double> k;
    return scalar_pair_pair(k, value, m1, s1, m2, s2);
}

void clr_vector_kernel_f32(rust::Slice<const float> in_vec,
                           float aux1_mean, float aux1_std,
                           rust::Slice<const float> aux2_means,
                           rust::Slice<const float> aux2_stds,
                           rust::Slice<float> out) {
    ::mcp::kernel::clr_vector_kernel<float, ::std::pair<float, float>, float> k;
    apply_pair_vector(k, in_vec, aux1_mean, aux1_std,
                      aux2_means, aux2_stds, out);
}
void clr_vector_kernel_f64(rust::Slice<const double> in_vec,
                           double aux1_mean, double aux1_std,
                           rust::Slice<const double> aux2_means,
                           rust::Slice<const double> aux2_stds,
                           rust::Slice<double> out) {
    ::mcp::kernel::clr_vector_kernel<double, ::std::pair<double, double>, double> k;
    apply_pair_vector(k, in_vec, aux1_mean, aux1_std,
                      aux2_means, aux2_stds, out);
}

void zscored_clr_kernel_f32(rust::Slice<const float> zi,
                            rust::Slice<const float> zj,
                            rust::Slice<float> out) {
    ::mcp::kernel::zscored_clr_kernel<float, float, float> k;
    apply_zscored_binary(k, zi, zj, out);
}
void zscored_clr_kernel_f64(rust::Slice<const double> zi,
                            rust::Slice<const double> zj,
                            rust::Slice<double> out) {
    ::mcp::kernel::zscored_clr_kernel<double, double, double> k;
    apply_zscored_binary(k, zi, zj, out);
}

// -----------------------------------------------------------------------------
// threshold.hpp
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline void threshold_impl(rust::Slice<const T> in_vec,
                           T min_thresh, T max_thresh, T default_val,
                           rust::Slice<T> out) {
    if (in_vec.size() == 0) return;
    assert(out.size() >= in_vec.size() &&
           "threshold_kernel: out buffer too small");
    ::mcp::kernel::threshold<T> kernel(min_thresh, max_thresh, default_val);
    kernel(in_vec.data(), in_vec.size(), out.data());
}

template <typename T>
inline void inverted_threshold_impl(rust::Slice<const T> in_vec,
                                    T min_thresh, T max_thresh, T default_val,
                                    rust::Slice<T> out) {
    if (in_vec.size() == 0) return;
    assert(out.size() >= in_vec.size() &&
           "inverted_threshold_kernel: out buffer too small");
    ::mcp::kernel::inverted_threshold<T> kernel(min_thresh, max_thresh, default_val);
    kernel(in_vec.data(), in_vec.size(), out.data());
}

template <typename T>
inline void threshold2_impl(rust::Slice<const T> in_vec,
                            rust::Slice<const T> aux,
                            T min_thresh, T max_thresh, T default_val,
                            rust::Slice<T> out) {
    const ::std::size_t count = in_vec.size();
    if (count == 0) return;
    assert(aux.size() >= count &&
           "threshold2_kernel: aux too short");
    assert(out.size() >= count &&
           "threshold2_kernel: out buffer too small");
    ::mcp::kernel::threshold2<T, T> kernel(min_thresh, max_thresh, default_val);
    kernel(in_vec.data(), aux.data(), count, out.data());
}

template <typename T>
inline void inverted_threshold2_impl(rust::Slice<const T> in_vec,
                                     rust::Slice<const T> aux,
                                     T min_thresh, T max_thresh, T default_val,
                                     rust::Slice<T> out) {
    const ::std::size_t count = in_vec.size();
    if (count == 0) return;
    assert(aux.size() >= count &&
           "inverted_threshold2_kernel: aux too short");
    assert(out.size() >= count &&
           "inverted_threshold2_kernel: out buffer too small");
    ::mcp::kernel::inverted_threshold2<T, T> kernel(min_thresh, max_thresh, default_val);
    kernel(in_vec.data(), aux.data(), count, out.data());
}

template <typename T>
inline void mask_impl(rust::Slice<const T> in_vec,
                      rust::Slice<const ::std::uint8_t> mask_vec,
                      bool negate, T default_val,
                      rust::Slice<T> out) {
    const ::std::size_t count = in_vec.size();
    if (count == 0) return;
    assert(mask_vec.size() >= count &&
           "mask_kernel: mask too short");
    assert(out.size() >= count &&
           "mask_kernel: out buffer too small");

    // Convert u8 mask to a temporary bool buffer to avoid relying on
    // bit-level ABI details of bool.
    ::std::vector<bool> bool_mask(count);
    ::std::vector<unsigned char> raw(count);
    for (::std::size_t i = 0; i < count; ++i) {
        raw[i] = mask_vec[i] != 0 ? 1 : 0;
    }
    // Run the kernel using a plain bool* array.
    ::std::vector<char> tmp(count); // char as 1-byte stand-in
    bool * mb = reinterpret_cast<bool *>(raw.data());
    if (negate) {
        ::mcp::kernel::mask<T, true> kernel(default_val);
        kernel(in_vec.data(), mb, count, out.data());
    } else {
        ::mcp::kernel::mask<T, false> kernel(default_val);
        kernel(in_vec.data(), mb, count, out.data());
    }
    (void)bool_mask;
    (void)tmp;
}

} // namespace

void threshold_kernel_f32(rust::Slice<const float> in_vec,
                          float mn, float mx, float def,
                          rust::Slice<float> out) {
    threshold_impl<float>(in_vec, mn, mx, def, out);
}
void threshold_kernel_f64(rust::Slice<const double> in_vec,
                          double mn, double mx, double def,
                          rust::Slice<double> out) {
    threshold_impl<double>(in_vec, mn, mx, def, out);
}

void inverted_threshold_kernel_f32(rust::Slice<const float> in_vec,
                                   float mn, float mx, float def,
                                   rust::Slice<float> out) {
    inverted_threshold_impl<float>(in_vec, mn, mx, def, out);
}
void inverted_threshold_kernel_f64(rust::Slice<const double> in_vec,
                                   double mn, double mx, double def,
                                   rust::Slice<double> out) {
    inverted_threshold_impl<double>(in_vec, mn, mx, def, out);
}

void threshold2_kernel_f32(rust::Slice<const float> in_vec,
                           rust::Slice<const float> aux,
                           float mn, float mx, float def,
                           rust::Slice<float> out) {
    threshold2_impl<float>(in_vec, aux, mn, mx, def, out);
}
void threshold2_kernel_f64(rust::Slice<const double> in_vec,
                           rust::Slice<const double> aux,
                           double mn, double mx, double def,
                           rust::Slice<double> out) {
    threshold2_impl<double>(in_vec, aux, mn, mx, def, out);
}

void inverted_threshold2_kernel_f32(rust::Slice<const float> in_vec,
                                    rust::Slice<const float> aux,
                                    float mn, float mx, float def,
                                    rust::Slice<float> out) {
    inverted_threshold2_impl<float>(in_vec, aux, mn, mx, def, out);
}
void inverted_threshold2_kernel_f64(rust::Slice<const double> in_vec,
                                    rust::Slice<const double> aux,
                                    double mn, double mx, double def,
                                    rust::Slice<double> out) {
    inverted_threshold2_impl<double>(in_vec, aux, mn, mx, def, out);
}

void mask_kernel_f32(rust::Slice<const float> in_vec,
                     rust::Slice<const ::std::uint8_t> mask_vec,
                     bool negate, float def,
                     rust::Slice<float> out) {
    mask_impl<float>(in_vec, mask_vec, negate, def, out);
}
void mask_kernel_f64(rust::Slice<const double> in_vec,
                     rust::Slice<const ::std::uint8_t> mask_vec,
                     bool negate, double def,
                     rust::Slice<double> out) {
    mask_impl<double>(in_vec, mask_vec, negate, def, out);
}

// -----------------------------------------------------------------------------
// dpi.hpp / mcp.hpp
// -----------------------------------------------------------------------------

namespace {

template <typename T>
inline ::std::vector<T> tfs_to_vec(rust::Slice<const T> tfs) {
    return ::std::vector<T>(tfs.data(), tfs.data() + tfs.size());
}

template <typename T>
inline void check_pair_rows(::std::size_t x, ::std::size_t z,
                            rust::Slice<const T> row_x,
                            rust::Slice<const T> row_z) {
    assert(row_x.size() == row_z.size() &&
           "row_x and row_z must have the same length");
    assert(x < row_x.size() && z < row_x.size() &&
           "x and z must be in range [0, count)");
    (void)x; (void)z;
}

template <typename T>
inline bool dpi_impl(::std::size_t x, ::std::size_t z,
                     rust::Slice<const T> row_x,
                     rust::Slice<const T> row_z,
                     double tolerance) {
    check_pair_rows<T>(x, z, row_x, row_z);
    ::mcp::kernel::dpi_kernel<T, true, true> kernel(tolerance);
    return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
}

template <typename T>
inline bool dpi_tf_impl(::std::size_t x, ::std::size_t z,
                        rust::Slice<const T> row_x,
                        rust::Slice<const T> row_z,
                        rust::Slice<const T> tfs,
                        double tolerance) {
    check_pair_rows<T>(x, z, row_x, row_z);
    assert(tfs.size() >= row_x.size() &&
           "dpi_tf_kernel: tfs must cover all nodes");
    auto tf_vec = tfs_to_vec<T>(tfs);
    ::mcp::kernel::dpi_tf_kernel<T, true> kernel(
        tf_vec, static_cast<T>(tolerance));
    return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
}

template <typename T>
inline T mcp2_maxmin_impl(::std::size_t x, ::std::size_t z,
                          rust::Slice<const T> row_x,
                          rust::Slice<const T> row_z,
                          rust::Slice<const T> tfs) {
    check_pair_rows<T>(x, z, row_x, row_z);
    if (tfs.size() == 0) {
        ::mcp::kernel::mcp2_maxmin_kernel<T, false, false> kernel;
        return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
    }
    assert(tfs.size() >= row_x.size() &&
           "mcp2_maxmin: tfs too short");
    auto tf_vec = tfs_to_vec<T>(tfs);
    ::mcp::kernel::mcp2_maxmin_kernel<T, true, false> kernel(tf_vec);
    return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
}

template <typename T>
inline T mcp_ratio_impl(::std::size_t x, ::std::size_t z,
                        rust::Slice<const T> row_x,
                        rust::Slice<const T> row_z,
                        rust::Slice<const T> tfs) {
    check_pair_rows<T>(x, z, row_x, row_z);
    if (tfs.size() == 0) {
        ::mcp::kernel::mcp_ratio_kernel<T, T, false> kernel;
        return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
    }
    assert(tfs.size() >= row_x.size() &&
           "mcp_ratio: tfs too short");
    auto tf_vec = tfs_to_vec<T>(tfs);
    ::mcp::kernel::mcp_ratio_kernel<T, T, true> kernel(tf_vec);
    return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
}

template <typename T>
inline T mcp_tolerance_impl(::std::size_t x, ::std::size_t z,
                            rust::Slice<const T> row_x,
                            rust::Slice<const T> row_z,
                            rust::Slice<const T> tfs,
                            bool clamped) {
    check_pair_rows<T>(x, z, row_x, row_z);
    if (tfs.size() == 0) {
        ::mcp::kernel::mcp_tolerance_kernel<T, T, false> kernel(clamped);
        return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
    }
    assert(tfs.size() >= row_x.size() &&
           "mcp_tolerance: tfs too short");
    auto tf_vec = tfs_to_vec<T>(tfs);
    ::mcp::kernel::mcp_tolerance_kernel<T, T, true> kernel(clamped, tf_vec);
    return kernel(x, z, row_x.data(), row_z.data(), row_x.size());
}

// NOTE: mcp::kernel::mcp_diff_kernel<float, float> fails to compile because
// the underlying header uses `std::max(0.0, out)` with `out` of type `float`.
// We replicate the (non-clamped) behavior directly here for both f32 and f64,
// since the kernel sets `clamped = std::is_unsigned<OT>::value`, which is
// always false for floating-point output types.
template <typename T>
inline T mcp_diff_impl(::std::size_t x, ::std::size_t z,
                       rust::Slice<const T> row_x,
                       rust::Slice<const T> row_z,
                       rust::Slice<const T> tfs) {
    check_pair_rows<T>(x, z, row_x, row_z);
    T mx = mcp2_maxmin_impl<T>(x, z, row_x, row_z, tfs);
    return row_x[z] - mx;
}

} // namespace

bool dpi_kernel_f32(::std::size_t x, ::std::size_t z,
                    rust::Slice<const float> row_x,
                    rust::Slice<const float> row_z,
                    double tolerance) {
    return dpi_impl<float>(x, z, row_x, row_z, tolerance);
}
bool dpi_kernel_f64(::std::size_t x, ::std::size_t z,
                    rust::Slice<const double> row_x,
                    rust::Slice<const double> row_z,
                    double tolerance) {
    return dpi_impl<double>(x, z, row_x, row_z, tolerance);
}

bool dpi_tf_kernel_f32(::std::size_t x, ::std::size_t z,
                       rust::Slice<const float> row_x,
                       rust::Slice<const float> row_z,
                       rust::Slice<const float> tfs,
                       double tolerance) {
    return dpi_tf_impl<float>(x, z, row_x, row_z, tfs, tolerance);
}
bool dpi_tf_kernel_f64(::std::size_t x, ::std::size_t z,
                       rust::Slice<const double> row_x,
                       rust::Slice<const double> row_z,
                       rust::Slice<const double> tfs,
                       double tolerance) {
    return dpi_tf_impl<double>(x, z, row_x, row_z, tfs, tolerance);
}

float mcp2_maxmin_kernel_f32(::std::size_t x, ::std::size_t z,
                             rust::Slice<const float> row_x,
                             rust::Slice<const float> row_z,
                             rust::Slice<const float> tfs) {
    return mcp2_maxmin_impl<float>(x, z, row_x, row_z, tfs);
}
double mcp2_maxmin_kernel_f64(::std::size_t x, ::std::size_t z,
                              rust::Slice<const double> row_x,
                              rust::Slice<const double> row_z,
                              rust::Slice<const double> tfs) {
    return mcp2_maxmin_impl<double>(x, z, row_x, row_z, tfs);
}

float mcp_ratio_kernel_f32(::std::size_t x, ::std::size_t z,
                           rust::Slice<const float> row_x,
                           rust::Slice<const float> row_z,
                           rust::Slice<const float> tfs) {
    return mcp_ratio_impl<float>(x, z, row_x, row_z, tfs);
}
double mcp_ratio_kernel_f64(::std::size_t x, ::std::size_t z,
                            rust::Slice<const double> row_x,
                            rust::Slice<const double> row_z,
                            rust::Slice<const double> tfs) {
    return mcp_ratio_impl<double>(x, z, row_x, row_z, tfs);
}

float mcp_tolerance_kernel_f32(::std::size_t x, ::std::size_t z,
                               rust::Slice<const float> row_x,
                               rust::Slice<const float> row_z,
                               rust::Slice<const float> tfs,
                               bool clamped) {
    return mcp_tolerance_impl<float>(x, z, row_x, row_z, tfs, clamped);
}
double mcp_tolerance_kernel_f64(::std::size_t x, ::std::size_t z,
                                rust::Slice<const double> row_x,
                                rust::Slice<const double> row_z,
                                rust::Slice<const double> tfs,
                                bool clamped) {
    return mcp_tolerance_impl<double>(x, z, row_x, row_z, tfs, clamped);
}

float mcp_diff_kernel_f32(::std::size_t x, ::std::size_t z,
                          rust::Slice<const float> row_x,
                          rust::Slice<const float> row_z,
                          rust::Slice<const float> tfs) {
    return mcp_diff_impl<float>(x, z, row_x, row_z, tfs);
}
double mcp_diff_kernel_f64(::std::size_t x, ::std::size_t z,
                           rust::Slice<const double> row_x,
                           rust::Slice<const double> row_z,
                           rust::Slice<const double> tfs) {
    return mcp_diff_impl<double>(x, z, row_x, row_z, tfs);
}