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

// DecisionTree.hpp — 決定木 (CART アルゴリズム)
//
// 分類 (Gini 不純度) と回帰 (MSE) の両方に対応。
// 特徴量: Matrix<T> の各行がサンプル、各列が特徴量。
// ラベル: 分類は std::vector<int>、回帰は std::vector<T>。

#ifndef CALX_ML_TREE_DECISIONTREE_HPP
#define CALX_ML_TREE_DECISIONTREE_HPP

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

namespace calx::ml {

    /// 決定木のタスク種別
    enum class TreeTask { Classification, Regression };

    namespace detail {

        // -----------------------------------------------------------------
        // ノード構造
        // -----------------------------------------------------------------
        template<typename T>
        struct TreeNode {
            bool isLeaf = false;

            // 内部ノード
            size_t featureIndex = 0;
            T threshold = T(0);
            std::unique_ptr<TreeNode> left;   // <= threshold
            std::unique_ptr<TreeNode> right;  // >  threshold

            // 葉ノード
            int classLabel = 0;              // 分類用
            T value = T(0);                  // 回帰用
            std::vector<T> classProbabilities; // 分類の確率分布
        };

        // -----------------------------------------------------------------
        // Gini 不純度 (分類)
        // -----------------------------------------------------------------
        inline double giniImpurity(const std::vector<int>& labels,
                                   const std::vector<size_t>& indices) {
            if (indices.empty()) return 0.0;
            std::unordered_map<int, size_t> counts;
            for (auto i : indices) counts[labels[i]]++;
            double impurity = 1.0;
            double n = static_cast<double>(indices.size());
            for (auto& [cls, cnt] : counts) {
                double p = static_cast<double>(cnt) / n;
                impurity -= p * p;
            }
            return impurity;
        }

        // -----------------------------------------------------------------
        // MSE (回帰)
        // -----------------------------------------------------------------
        template<typename T>
        double mseImpurity(const std::vector<T>& values,
                           const std::vector<size_t>& indices) {
            if (indices.empty()) return 0.0;
            double mean = 0.0;
            for (auto i : indices) mean += static_cast<double>(values[i]);
            mean /= static_cast<double>(indices.size());
            double mse = 0.0;
            for (auto i : indices) {
                double d = static_cast<double>(values[i]) - mean;
                mse += d * d;
            }
            return mse / static_cast<double>(indices.size());
        }

        // -----------------------------------------------------------------
        // 最良分割の探索
        // -----------------------------------------------------------------
        template<typename T>
        struct SplitResult {
            bool found = false;
            size_t featureIndex = 0;
            T threshold = T(0);
            std::vector<size_t> leftIndices;
            std::vector<size_t> rightIndices;
            double gain = 0.0;
        };

        template<typename T>
        SplitResult<T> findBestSplitClassification(
            const Matrix<T>& X, const std::vector<int>& y,
            const std::vector<size_t>& indices, size_t numFeatures)
        {
            SplitResult<T> best;
            double parentImpurity = giniImpurity(y, indices);
            double n = static_cast<double>(indices.size());

            // 特徴量ごとの値をソートして分割候補を探索
            std::vector<std::pair<T, size_t>> featureValues(indices.size());

            for (size_t f = 0; f < numFeatures; ++f) {
                for (size_t k = 0; k < indices.size(); ++k)
                    featureValues[k] = { X(indices[k], f), indices[k] };
                std::sort(featureValues.begin(), featureValues.end(),
                          [](const auto& a, const auto& b) { return a.first < b.first; });

                // 隣接サンプル間の中点を分割候補とする
                for (size_t k = 0; k + 1 < featureValues.size(); ++k) {
                    if (featureValues[k].first == featureValues[k + 1].first)
                        continue;
                    T thresh = (featureValues[k].first + featureValues[k + 1].first) / T(2);

                    std::vector<size_t> leftIdx, rightIdx;
                    leftIdx.reserve(k + 1);
                    rightIdx.reserve(featureValues.size() - k - 1);
                    for (size_t j = 0; j <= k; ++j)
                        leftIdx.push_back(featureValues[j].second);
                    for (size_t j = k + 1; j < featureValues.size(); ++j)
                        rightIdx.push_back(featureValues[j].second);

                    double leftImp = giniImpurity(y, leftIdx);
                    double rightImp = giniImpurity(y, rightIdx);
                    double wLeft = static_cast<double>(leftIdx.size()) / n;
                    double wRight = static_cast<double>(rightIdx.size()) / n;
                    double gain = parentImpurity - wLeft * leftImp - wRight * rightImp;

                    if (gain > best.gain) {
                        best.found = true;
                        best.featureIndex = f;
                        best.threshold = thresh;
                        best.leftIndices = std::move(leftIdx);
                        best.rightIndices = std::move(rightIdx);
                        best.gain = gain;
                    }
                }
            }
            return best;
        }

        template<typename T>
        SplitResult<T> findBestSplitRegression(
            const Matrix<T>& X, const std::vector<T>& y,
            const std::vector<size_t>& indices, size_t numFeatures)
        {
            SplitResult<T> best;
            double parentMSE = mseImpurity(y, indices);
            double n = static_cast<double>(indices.size());

            std::vector<std::pair<T, size_t>> featureValues(indices.size());

            for (size_t f = 0; f < numFeatures; ++f) {
                for (size_t k = 0; k < indices.size(); ++k)
                    featureValues[k] = { X(indices[k], f), indices[k] };
                std::sort(featureValues.begin(), featureValues.end(),
                          [](const auto& a, const auto& b) { return a.first < b.first; });

                for (size_t k = 0; k + 1 < featureValues.size(); ++k) {
                    if (featureValues[k].first == featureValues[k + 1].first)
                        continue;
                    T thresh = (featureValues[k].first + featureValues[k + 1].first) / T(2);

                    std::vector<size_t> leftIdx, rightIdx;
                    leftIdx.reserve(k + 1);
                    rightIdx.reserve(featureValues.size() - k - 1);
                    for (size_t j = 0; j <= k; ++j)
                        leftIdx.push_back(featureValues[j].second);
                    for (size_t j = k + 1; j < featureValues.size(); ++j)
                        rightIdx.push_back(featureValues[j].second);

                    double leftMSE = mseImpurity(y, leftIdx);
                    double rightMSE = mseImpurity(y, rightIdx);
                    double wLeft = static_cast<double>(leftIdx.size()) / n;
                    double wRight = static_cast<double>(rightIdx.size()) / n;
                    double gain = parentMSE - wLeft * leftMSE - wRight * rightMSE;

                    if (gain > best.gain) {
                        best.found = true;
                        best.featureIndex = f;
                        best.threshold = thresh;
                        best.leftIndices = std::move(leftIdx);
                        best.rightIndices = std::move(rightIdx);
                        best.gain = gain;
                    }
                }
            }
            return best;
        }

    } // namespace detail

    // =====================================================================
    // DecisionTreeClassifier
    // =====================================================================

    /// CART 分類木
    template<typename T>
    class DecisionTreeClassifier {
    public:
        /// @param maxDepth    最大深さ (0 = 無制限)
        /// @param minSamples  分割に必要な最小サンプル数
        explicit DecisionTreeClassifier(size_t maxDepth = 0, size_t minSamples = 2)
            : maxDepth_(maxDepth), minSamples_(minSamples) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<int>& y) {
            assert(X.rows() == y.size());
            numFeatures_ = X.cols();
            std::vector<size_t> indices(y.size());
            std::iota(indices.begin(), indices.end(), 0);

            // クラスラベル一覧を収集
            classes_.clear();
            for (int label : y) {
                if (std::find(classes_.begin(), classes_.end(), label) == classes_.end())
                    classes_.push_back(label);
            }
            std::sort(classes_.begin(), classes_.end());

            root_ = buildClassification(X, y, indices, 1);
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            assert(root_);
            const auto* node = root_.get();
            while (!node->isLeaf) {
                if (X(row, node->featureIndex) <= node->threshold)
                    node = node->left.get();
                else
                    node = node->right.get();
            }
            return node->classLabel;
        }

        /// 全サンプルの予測
        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;
        }

        /// 単一サンプルのクラス確率
        std::vector<T> predictProbability(const Matrix<T>& X, size_t row) const {
            assert(root_);
            const auto* node = root_.get();
            while (!node->isLeaf) {
                if (X(row, node->featureIndex) <= node->threshold)
                    node = node->left.get();
                else
                    node = node->right.get();
            }
            return node->classProbabilities;
        }

        /// 木の深さ
        size_t depth() const { return root_ ? nodeDepth(root_.get()) : 0; }

        /// 葉ノード数
        size_t numLeaves() const { return root_ ? countLeaves(root_.get()) : 0; }

        /// クラスラベル一覧
        const std::vector<int>& classes() const { return classes_; }

    private:
        size_t maxDepth_;
        size_t minSamples_;
        size_t numFeatures_ = 0;
        std::vector<int> classes_;
        std::unique_ptr<detail::TreeNode<T>> root_;

        std::unique_ptr<detail::TreeNode<T>> buildClassification(
            const Matrix<T>& X, const std::vector<int>& y,
            const std::vector<size_t>& indices, size_t currentDepth)
        {
            auto node = std::make_unique<detail::TreeNode<T>>();

            // 葉ノード条件: サンプル数不足 / 深さ上限 / 純粋
            bool allSame = std::all_of(indices.begin(), indices.end(),
                [&](size_t i) { return y[i] == y[indices[0]]; });

            if (allSame || indices.size() < minSamples_ ||
                (maxDepth_ > 0 && currentDepth >= maxDepth_))
            {
                makeClassificationLeaf(*node, y, indices);
                return node;
            }

            auto split = detail::findBestSplitClassification(X, y, indices, numFeatures_);
            if (!split.found) {
                makeClassificationLeaf(*node, y, indices);
                return node;
            }

            node->featureIndex = split.featureIndex;
            node->threshold = split.threshold;
            node->left  = buildClassification(X, y, split.leftIndices,  currentDepth + 1);
            node->right = buildClassification(X, y, split.rightIndices, currentDepth + 1);
            return node;
        }

        void makeClassificationLeaf(detail::TreeNode<T>& node,
                                     const std::vector<int>& y,
                                     const std::vector<size_t>& indices) const
        {
            node.isLeaf = true;

            // 多数決でクラスラベル決定
            std::unordered_map<int, size_t> counts;
            for (auto i : indices) counts[y[i]]++;
            int bestLabel = y[indices[0]];
            size_t bestCount = 0;
            for (auto& [label, cnt] : counts) {
                if (cnt > bestCount) { bestCount = cnt; bestLabel = label; }
            }
            node.classLabel = bestLabel;

            // クラス確率
            node.classProbabilities.resize(classes_.size(), T(0));
            double n = static_cast<double>(indices.size());
            for (size_t c = 0; c < classes_.size(); ++c) {
                auto it = counts.find(classes_[c]);
                if (it != counts.end())
                    node.classProbabilities[c] = static_cast<T>(it->second / n);
            }
        }

        static size_t nodeDepth(const detail::TreeNode<T>* node) {
            if (!node || node->isLeaf) return 1;
            return 1 + std::max(nodeDepth(node->left.get()),
                                nodeDepth(node->right.get()));
        }

        static size_t countLeaves(const detail::TreeNode<T>* node) {
            if (!node) return 0;
            if (node->isLeaf) return 1;
            return countLeaves(node->left.get()) + countLeaves(node->right.get());
        }
    };

    // =====================================================================
    // DecisionTreeRegressor
    // =====================================================================

    /// CART 回帰木
    template<typename T>
    class DecisionTreeRegressor {
    public:
        /// @param maxDepth    最大深さ (0 = 無制限)
        /// @param minSamples  分割に必要な最小サンプル数
        explicit DecisionTreeRegressor(size_t maxDepth = 0, size_t minSamples = 2)
            : maxDepth_(maxDepth), minSamples_(minSamples) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<T>& y) {
            assert(X.rows() == y.size());
            numFeatures_ = X.cols();
            std::vector<size_t> indices(y.size());
            std::iota(indices.begin(), indices.end(), 0);
            root_ = buildRegression(X, y, indices, 1);
        }

        /// 単一サンプルの予測
        T predict(const Matrix<T>& X, size_t row) const {
            assert(root_);
            const auto* node = root_.get();
            while (!node->isLeaf) {
                if (X(row, node->featureIndex) <= node->threshold)
                    node = node->left.get();
                else
                    node = node->right.get();
            }
            return node->value;
        }

        /// 全サンプルの予測
        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;
        }

        /// 木の深さ
        size_t depth() const { return root_ ? nodeDepth(root_.get()) : 0; }

        /// 葉ノード数
        size_t numLeaves() const { return root_ ? countLeaves(root_.get()) : 0; }

    private:
        size_t maxDepth_;
        size_t minSamples_;
        size_t numFeatures_ = 0;
        std::unique_ptr<detail::TreeNode<T>> root_;

        std::unique_ptr<detail::TreeNode<T>> buildRegression(
            const Matrix<T>& X, const std::vector<T>& y,
            const std::vector<size_t>& indices, size_t currentDepth)
        {
            auto node = std::make_unique<detail::TreeNode<T>>();

            if (indices.size() < minSamples_ ||
                (maxDepth_ > 0 && currentDepth >= maxDepth_))
            {
                makeRegressionLeaf(*node, y, indices);
                return node;
            }

            // 全サンプルが同一値なら葉
            bool allSame = std::all_of(indices.begin(), indices.end(),
                [&](size_t i) { return y[i] == y[indices[0]]; });
            if (allSame) {
                makeRegressionLeaf(*node, y, indices);
                return node;
            }

            auto split = detail::findBestSplitRegression(X, y, indices, numFeatures_);
            if (!split.found) {
                makeRegressionLeaf(*node, y, indices);
                return node;
            }

            node->featureIndex = split.featureIndex;
            node->threshold = split.threshold;
            node->left  = buildRegression(X, y, split.leftIndices,  currentDepth + 1);
            node->right = buildRegression(X, y, split.rightIndices, currentDepth + 1);
            return node;
        }

        static void makeRegressionLeaf(detail::TreeNode<T>& node,
                                         const std::vector<T>& y,
                                         const std::vector<size_t>& indices)
        {
            node.isLeaf = true;
            T sum = T(0);
            for (auto i : indices) sum += y[i];
            node.value = sum / static_cast<T>(indices.size());
        }

        static size_t nodeDepth(const detail::TreeNode<T>* node) {
            if (!node || node->isLeaf) return 1;
            return 1 + std::max(nodeDepth(node->left.get()),
                                nodeDepth(node->right.get()));
        }

        static size_t countLeaves(const detail::TreeNode<T>* node) {
            if (!node) return 0;
            if (node->isLeaf) return 1;
            return countLeaves(node->left.get()) + countLeaves(node->right.get());
        }
    };

} // namespace calx::ml

#endif // CALX_ML_TREE_DECISIONTREE_HPP
