kev/Drawer/Module/GeoSigmaDraw/GridAlgoInterpolator.h

#pragma once
#include <vector>
#include <cmath>
#include <algorithm>
#include <limits>
#include <string>
#include <random>
#include <iostream>
#include <cassert>
#include <unordered_set>
#include <omp.h>

#ifndef NOMINMAX
#define NOMINMAX
#endif

#if defined(min)
#undef min
#endif

#if defined(max)
#undef max
#endif
#include "DrawOperator/nanoflann.hpp" 

namespace NItem
{
    // ---------------------------------------------------------
    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ݽṹ
    // ---------------------------------------------------------
    struct AlgoPoint {
        double x, y, z;
    };

    // ---------------------------------------------------------
    // [<5B><><EFBFBD><EFBFBD>] <20><><EFBFBD><EFBFBD><EFBFBD>μ<EFBFBD><CEBC><EFBFBD><EFBFBD>㷨<EFBFBD><E3B7A8><EFBFBD><EFBFBD>
    // ---------------------------------------------------------
    class CGeoAlgo {
    public:

        // [<5B><><EFBFBD><EFBFBD>] һά<D2BB><CEAC><EFBFBD><EFBFBD> (<28><><EFBFBD><EFBFBD>ʵ<EFBFBD><CAB5><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>)
        static void Dilate1D(const std::vector<unsigned char>& in, std::vector<unsigned char>& out,
            int width, int height, int radius, bool horizontal) {

            // <20><><EFBFBD><EFBFBD>ѭ<EFBFBD><D1AD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ÿһ<C3BF><D2BB>(<28><><EFBFBD><EFBFBD>)
            int major_dim = horizontal ? height : width;
            int minor_dim = horizontal ? width : height;

            // OpenMP <20><><EFBFBD><EFBFBD>
#pragma omp parallel for
            for (int i = 0; i < major_dim; ++i) {
                for (int j = 0; j < minor_dim; ++j) {
                    int current_idx = horizontal ? (i * width + j) : (j * width + i);

                    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǰ<EFBFBD><C7B0><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ч<EFBFBD><D0A7>Ĥ
                    if (in[current_idx] != 0) {
                        // <20><><EFBFBD><EFBFBD>Ӱ<EFBFBD>췶Χ [start, end]
                        int start = std::max(0, j - radius);
                        int end = std::min(minor_dim - 1, j + radius);

                        for (int k = start; k <= end; ++k) {
                            int target_idx = horizontal ? (i * width + k) : (k * width + i);
                            out[target_idx] = 1;
                        }
                    }
                }
            }
        }
    };

    // ---------------------------------------------------------
    // [<5B>Զ<EFBFBD><D4B6><EFBFBD>] <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
    // ---------------------------------------------------------
    struct SimpleMatrix {
        int rows, cols;
        std::vector<double> data;
        SimpleMatrix(int r, int c) : rows(r), cols(c), data(r* c, 0.0) {}
        inline double& operator()(int r, int c) { return data[r * cols + c]; }
        inline const double& operator()(int r, int c) const { return data[r * cols + c]; }
        void setZero() { std::fill(data.begin(), data.end(), 0.0); }
    };

    // ---------------------------------------------------------
    // [<5B>Զ<EFBFBD><D4B6><EFBFBD>] <20><>˹<EFBFBD><CBB9>Ԫ<EFBFBD><D4AA><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Է<EFBFBD><D4B7><EFBFBD><EFBFBD><EFBFBD>
    // ---------------------------------------------------------
    inline bool SolveLinearSystem(SimpleMatrix A, std::vector<double> B, std::vector<double>& X) {
        int n = A.rows;
        X.resize(n);
        for (int i = 0; i < n; ++i) {
            int pivot = i;
            double maxVal = std::abs(A(i, i));
            for (int j = i + 1; j < n; ++j) {
                if (std::abs(A(j, i)) > maxVal) { maxVal = std::abs(A(j, i)); pivot = j; }
            }
            if (maxVal < 1e-14) return false;
            if (pivot != i) {
                for (int k = i; k < n; ++k) std::swap(A(i, k), A(pivot, k));
                std::swap(B[i], B[pivot]);
            }
            for (int j = i + 1; j < n; ++j) {
                double factor = A(j, i) / A(i, i);
                for (int k = i + 1; k < n; ++k) A(j, k) -= factor * A(i, k);
                B[j] -= factor * B[i];
            }
        }
        for (int i = n - 1; i >= 0; --i) {
            double sum = B[i];
            for (int j = i + 1; j < n; ++j) sum -= A(i, j) * X[j];
            X[i] = sum / A(i, i);
        }
        return true;
    }

    // ---------------------------------------------------------
    // [<5B>޸<EFBFBD>] KernelType ö<><C3B6>
    // <20>Ƴ<EFBFBD><C6B3><EFBFBD> K_THIN_PLATE <20><> K_CUBIC
    // ---------------------------------------------------------
    enum KernelType {
        K_LINEAR,
        K_MULTIQUIDRIC,
        K_INV_MULTIQUADRIC,
        K_GAUSSIAN
    };

    struct RbfParams {
        KernelType kernel = K_LINEAR; //  Ĭ<><C4AC>ֵ<EFBFBD><D6B5>Ϊ K_LINEAR
        double epsilon = 0.0;
        double smoothing = 3.0;
        int neighbors = 60;
        double max_dist = -1.0;
        bool auto_epsilon = true;
        int eps_seed = 1234;
    };


    // ---------------------------------------------------------
    // nanoflann KDTree <20><><EFBFBD><EFBFBD>
    // ---------------------------------------------------------
    using kd_index_t = size_t;

    struct PointCloudAdaptor {
        const std::vector<AlgoPoint>& pts;
        PointCloudAdaptor(const std::vector<AlgoPoint>& p) : pts(p) {}
        inline size_t kdtree_get_point_count() const { return pts.size(); }
        inline double kdtree_get_pt(const size_t idx, const size_t dim) const {
            return (dim == 0) ? pts[idx].x : pts[idx].y;
        }
        template <class BBOX> bool kdtree_get_bbox(BBOX& /*bb*/) const { return false; }
    };

    typedef nanoflann::KDTreeSingleIndexAdaptor<
        nanoflann::L2_Simple_Adaptor<double, PointCloudAdaptor>,
        PointCloudAdaptor,
        2,
        kd_index_t
    > MyKDTree;


    // ---------------------------------------------------------
    // RBF <20><>ֵ<EFBFBD><D6B5><EFBFBD><EFBFBD>
    // ---------------------------------------------------------
    class RBFInterpolatorCpp {
    public:
        RBFInterpolatorCpp(const std::vector<AlgoPoint>& points)
            : pts(points), adaptor(pts), kdtree(2, adaptor, nanoflann::KDTreeSingleIndexAdaptorParams(10))
        {
            if (pts.empty()) throw std::runtime_error("RBF Error: Input points empty");
            kdtree.buildIndex();
        }

        double get_nearest_value(double x, double y) const {
            double query_pt[2] = { x, y };
            kd_index_t out_idx; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
            double out_d2;      // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ƽ<EFBFBD><C6BD>

            // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> 1 <20><><EFBFBD><EFBFBD>
            size_t num_matches = kdtree.knnSearch(&query_pt[0], 1, &out_idx, &out_d2);

            if (num_matches > 0) {
                return pts[out_idx].z;
            }
            return std::numeric_limits<double>::quiet_NaN();
        }

        double auto_epsilon(double max_dist, int rng_seed = 1234) const {
            if (max_dist > 0.0) return std::max(max_dist / 3.0, 1e-9);
            size_t N = pts.size();
            size_t m = std::min<size_t>(2000, N);
            std::vector<size_t> idxs;
            idxs.reserve(m);
            std::mt19937_64 rng((uint64_t)rng_seed);
            if (m == N) { for (size_t i = 0; i < m; ++i) idxs.push_back(i); }
            else {
                std::uniform_int_distribution<size_t> dist(0, N - 1);
                std::unordered_set<size_t> used;
                while (used.size() < m) used.insert(dist(rng));
                for (auto v : used) idxs.push_back(v);
            }
            std::vector<kd_index_t> out_idx(2);
            std::vector<double> out_d2(2);
            std::vector<double> nn;
            for (size_t i : idxs) {
                double query_pt[2] = { pts[i].x, pts[i].y };
                if (kdtree.knnSearch(&query_pt[0], 2, &out_idx[0], &out_d2[0]) >= 2) {
                    double d = std::sqrt(out_d2[1]);
                    if (d > 1e-12) nn.push_back(d);
                }
            }
            if (nn.empty()) return 1.0;
            size_t mid = nn.size() / 2;
            std::nth_element(nn.begin(), nn.begin() + mid, nn.end());
            return std::max(nn[mid] * 2.0, 1e-9);
        }

        static double kernel_value(double r, KernelType k, double epsilon) {
            if (r <= 0.0) {
                return (k == K_GAUSSIAN || k == K_INV_MULTIQUADRIC || k == K_MULTIQUIDRIC) ? 1.0 : 0.0;
            }
            switch (k) {
            case K_LINEAR: return r;
            case K_MULTIQUIDRIC:
                if (epsilon <= 0) epsilon = 1.0;
                return std::sqrt(1.0 + (r * r) / (epsilon * epsilon));
            case K_INV_MULTIQUADRIC:
                if (epsilon <= 0) epsilon = 1.0;
                return 1.0 / std::sqrt(1.0 + (r * r) / (epsilon * epsilon));
            case K_GAUSSIAN:
                if (epsilon <= 0) epsilon = 1.0;
                return std::exp(-(r * r) / (epsilon * epsilon));
            default: return r;
            }
        }

        std::vector<double> predict(const std::vector<double>& xq, const std::vector<double>& yq,
            const RbfParams& params, int omp_threads = 1, int chunk = 40000) const
        {
            if (xq.size() != yq.size()) throw std::runtime_error("RBF Error: xq/yq size mismatch");
            size_t Q = xq.size();
            std::vector<double> out(Q, std::numeric_limits<double>::quiet_NaN());
            if (pts.empty()) return out;

            double eps_use = params.epsilon;
            bool need_eps = (params.kernel == K_GAUSSIAN || params.kernel == K_MULTIQUIDRIC || params.kernel == K_INV_MULTIQUADRIC);
            if (need_eps && (eps_use <= 0.0) && params.auto_epsilon) {
                eps_use = auto_epsilon(params.max_dist, params.eps_seed);
            }

            std::vector<char> ok_mask(Q, 1);
            if (params.max_dist > 0.0) {
                double md = params.max_dist;
                for (size_t i = 0; i < Q; ++i) {
                    double qpt[2] = { xq[i], yq[i] };
                    kd_index_t tmp_idx;
                    double tmp_d2;
                    if (kdtree.knnSearch(&qpt[0], 1, &tmp_idx, &tmp_d2) == 0 || std::sqrt(tmp_d2) > md) ok_mask[i] = 0;
                }
            }

            const size_t neighbors = std::max(5, std::min((int)params.neighbors, (int)pts.size()));
#ifdef _OPENMP
            if (omp_threads > 0) omp_set_num_threads(omp_threads);
#endif

            for (size_t s = 0; s < Q; s += chunk) {
                size_t e = std::min(Q, s + chunk);

#pragma omp parallel for
                for (int qi = (int)s; qi < (int)e; ++qi) {
                    if (!ok_mask[qi]) continue;

                    double query_pt[2] = { xq[qi], yq[qi] };
                    std::vector<kd_index_t> out_idx(neighbors);
                    std::vector<double> out_d2(neighbors);
                    size_t got = kdtree.knnSearch(&query_pt[0], neighbors, out_idx.data(), out_d2.data());
                    if (got == 0) continue;

                    double maxd2 = (params.max_dist > 0.0) ? (params.max_dist * params.max_dist) : std::numeric_limits<double>::infinity();
                    std::vector<const AlgoPoint*> valid;
                    valid.reserve(got);
                    for (size_t i = 0; i < got; ++i) {
                        if (out_d2[i] <= maxd2 + 1e-12) valid.push_back(&pts[out_idx[i]]);
                    }

                    int n = (int)valid.size();
                    if (n < 3) continue;

                    int m = n + 3;
                    SimpleMatrix A(m, m);
                    std::vector<double> B(m, 0.0);

                    for (int i = 0; i < n; ++i) {
                        B[i] = valid[i]->z;
                        for (int j = 0; j < n; ++j) {
                            double dx = valid[i]->x - valid[j]->x;
                            double dy = valid[i]->y - valid[j]->y;
                            double d = std::sqrt(dx * dx + dy * dy);
                            A(i, j) = kernel_value(d, params.kernel, eps_use);
                        }
                        A(i, i) += params.smoothing;
                    }
                    for (int i = 0; i < n; ++i) {
                        A(i, n + 0) = 1.0;
                        A(i, n + 1) = valid[i]->x;
                        A(i, n + 2) = valid[i]->y;
                        A(n + 0, i) = 1.0;
                        A(n + 1, i) = valid[i]->x;
                        A(n + 2, i) = valid[i]->y;
                    }

                    std::vector<double> X;
                    bool solved = false;
                    double jitter = 0.0;
                    double traceA = 0.0;
                    for (int i = 0; i < n; ++i) traceA += std::abs(A(i, i));
                    double base_jitter = std::max(1e-12, traceA * 1e-12);
                    int jitter_attempts = 0;
                    while (!solved && jitter_attempts < 6) {
                        //if (jitter > 0.0) { for (int i = 0; i < n; ++i) A(i, i) += jitter; }
                        if (jitter > 0.0) {
                            for (int i = 0; i < m; ++i) A(i, i) += jitter;
                        }
                        solved = SolveLinearSystem(A, B, X);
                        if (solved) break;
                        jitter_attempts++;
                        //jitter = (jitter == 0.0) ? base_jitter : jitter * 10.0;
                        double actual_jitter = (jitter == 0.0) ? std::max(base_jitter, 1e-10) : jitter * 10.0;
                        jitter = actual_jitter;
                    }

                    if (!solved) continue;

                    double val = 0.0;
                    for (int i = 0; i < n; ++i) {
                        double dx = xq[qi] - valid[i]->x;
                        double dy = yq[qi] - valid[i]->y;
                        double d = std::sqrt(dx * dx + dy * dy);
                        val += X[i] * kernel_value(d, params.kernel, eps_use);
                    }
                    val += X[n + 0] + X[n + 1] * xq[qi] + X[n + 2] * yq[qi];
                    out[qi] = val;
                }
            }
            return out;
        }
        size_t point_count() const { return pts.size(); }
    private:
        std::vector<AlgoPoint> pts;
        PointCloudAdaptor adaptor;
        MyKDTree kdtree;
    };
}