// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// KNN.hpp — k 近傍法
//
// 分類 (多数決) と回帰 (平均) の両方に対応。
// kd-tree による高速近傍探索 (低次元向け)。
// 高次元ではブルートフォースにフォールバック。

#ifndef CALX_ML_NEIGHBORS_KNN_HPP
#define CALX_ML_NEIGHBORS_KNN_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <functional>
#include <limits>
#include <numeric>
#include <queue>
#include <unordered_map>
#include <vector>

namespace calx::ml {

    namespace detail {

        // =================================================================
        // kd-tree
        // =================================================================
        template<typename T>
        struct KdNode {
            size_t index;           // 元データの行インデックス
            int splitDim;           // 分割次元
            T splitVal;             // 分割値
            std::unique_ptr<KdNode> left;
            std::unique_ptr<KdNode> right;
        };

        template<typename T>
        class KdTree {
        public:
            void build(const Matrix<T>& X) {
                data_ = &X;
                dim_ = X.cols();
                std::vector<size_t> indices(X.rows());
                std::iota(indices.begin(), indices.end(), 0);
                root_ = buildNode(indices, 0);
            }

            /// k 近傍を探索 (距離の昇順)
            /// 結果: (距離², インデックス) のペア
            std::vector<std::pair<T, size_t>> kNearest(
                const Matrix<T>& query, size_t row, size_t k) const
            {
                // 最大ヒープで k 個を管理
                using Pair = std::pair<T, size_t>;
                std::priority_queue<Pair> heap;
                searchNode(root_.get(), query, row, k, heap);

                std::vector<Pair> result;
                result.reserve(heap.size());
                while (!heap.empty()) {
                    result.push_back(heap.top());
                    heap.pop();
                }
                std::reverse(result.begin(), result.end());
                return result;
            }

        private:
            const Matrix<T>* data_ = nullptr;
            size_t dim_ = 0;
            std::unique_ptr<KdNode<T>> root_;

            std::unique_ptr<KdNode<T>> buildNode(std::vector<size_t>& indices, int depth) {
                if (indices.empty()) return nullptr;

                int axis = depth % static_cast<int>(dim_);

                // 中央値で分割
                std::sort(indices.begin(), indices.end(),
                    [&](size_t a, size_t b) {
                        return (*data_)(a, axis) < (*data_)(b, axis);
                    });

                size_t mid = indices.size() / 2;
                auto node = std::make_unique<KdNode<T>>();
                node->index = indices[mid];
                node->splitDim = axis;
                node->splitVal = (*data_)(indices[mid], axis);

                std::vector<size_t> leftIdx(indices.begin(), indices.begin() + mid);
                std::vector<size_t> rightIdx(indices.begin() + mid + 1, indices.end());

                node->left = buildNode(leftIdx, depth + 1);
                node->right = buildNode(rightIdx, depth + 1);
                return node;
            }

            T squaredDist(const Matrix<T>& query, size_t qRow, size_t dataRow) const {
                T dist = T(0);
                for (size_t j = 0; j < dim_; ++j) {
                    T d = query(qRow, j) - (*data_)(dataRow, j);
                    dist += d * d;
                }
                return dist;
            }

            void searchNode(const KdNode<T>* node,
                            const Matrix<T>& query, size_t qRow, size_t k,
                            std::priority_queue<std::pair<T, size_t>>& heap) const
            {
                if (!node) return;

                T dist = squaredDist(query, qRow, node->index);
                if (heap.size() < k) {
                    heap.push({ dist, node->index });
                } else if (dist < heap.top().first) {
                    heap.pop();
                    heap.push({ dist, node->index });
                }

                T diff = query(qRow, node->splitDim) - node->splitVal;
                KdNode<T>* first = diff <= T(0) ? node->left.get() : node->right.get();
                KdNode<T>* second = diff <= T(0) ? node->right.get() : node->left.get();

                searchNode(first, query, qRow, k, heap);

                // 超平面との距離が現在の k 番目より近ければ反対側も探索
                T planeDist = diff * diff;
                if (heap.size() < k || planeDist < heap.top().first)
                    searchNode(second, query, qRow, k, heap);
            }
        };

    } // namespace detail

    // =====================================================================
    // KNNClassifier
    // =====================================================================

    /// k 近傍分類器
    template<typename T>
    class KNNClassifier {
    public:
        /// @param k  近傍数
        explicit KNNClassifier(size_t k = 5) : k_(k) {}

        /// 学習 (データを記憶 + kd-tree 構築)
        void fit(const Matrix<T>& X, const std::vector<int>& y) {
            assert(X.rows() == y.size());
            X_ = X;
            y_ = y;
            tree_.build(X_);
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            auto neighbors = tree_.kNearest(X, row, k_);
            std::unordered_map<int, size_t> votes;
            for (auto& [dist, idx] : neighbors)
                votes[y_[idx]]++;
            int bestLabel = 0;
            size_t bestCount = 0;
            for (auto& [label, cnt] : votes) {
                if (cnt > bestCount) { bestCount = cnt; bestLabel = label; }
            }
            return bestLabel;
        }

        /// 全サンプルの予測
        std::vector<int> predict(const Matrix<T>& X) const {
            std::vector<int> result(X.rows());
            for (size_t i = 0; i < X.rows(); ++i)
                result[i] = predict(X, i);
            return result;
        }

    private:
        size_t k_;
        Matrix<T> X_;
        std::vector<int> y_;
        detail::KdTree<T> tree_;
    };

    // =====================================================================
    // KNNRegressor
    // =====================================================================

    /// k 近傍回帰器
    template<typename T>
    class KNNRegressor {
    public:
        /// @param k  近傍数
        explicit KNNRegressor(size_t k = 5) : k_(k) {}

        /// 学習 (データを記憶 + kd-tree 構築)
        void fit(const Matrix<T>& X, const std::vector<T>& y) {
            assert(X.rows() == y.size());
            X_ = X;
            y_ = y;
            tree_.build(X_);
        }

        /// 単一サンプルの予測
        T predict(const Matrix<T>& X, size_t row) const {
            auto neighbors = tree_.kNearest(X, row, k_);
            T sum = T(0);
            for (auto& [dist, idx] : neighbors)
                sum += y_[idx];
            return sum / static_cast<T>(neighbors.size());
        }

        /// 全サンプルの予測
        std::vector<T> predict(const Matrix<T>& X) const {
            std::vector<T> result(X.rows());
            for (size_t i = 0; i < X.rows(); ++i)
                result[i] = predict(X, i);
            return result;
        }

    private:
        size_t k_;
        Matrix<T> X_;
        std::vector<T> y_;
        detail::KdTree<T> tree_;
    };

} // namespace calx::ml

#endif // CALX_ML_NEIGHBORS_KNN_HPP
