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

// RandomForest.hpp — ランダムフォレスト
//
// Bagging + 特徴量サブサンプリング + CART 決定木のアンサンブル。
// 分類 (多数決) と回帰 (平均) の両方に対応。

#ifndef CALX_ML_ENSEMBLE_RANDOMFOREST_HPP
#define CALX_ML_ENSEMBLE_RANDOMFOREST_HPP

#include <ml/tree/DecisionTree.hpp>
#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <numeric>
#include <random>
#include <unordered_map>
#include <vector>

namespace calx::ml {

    // =====================================================================
    // RandomForestClassifier
    // =====================================================================

    /// ランダムフォレスト分類器
    template<typename T>
    class RandomForestClassifier {
    public:
        /// @param numTrees      木の数
        /// @param maxDepth      各木の最大深さ (0 = 無制限)
        /// @param minSamples    分割に必要な最小サンプル数
        /// @param maxFeatures   各木で使う特徴量数 (0 = sqrt(全特徴量))
        /// @param seed          乱数シード
        explicit RandomForestClassifier(size_t numTrees = 100,
                                         size_t maxDepth = 0,
                                         size_t minSamples = 2,
                                         size_t maxFeatures = 0,
                                         unsigned seed = 42)
            : numTrees_(numTrees), maxDepth_(maxDepth),
              minSamples_(minSamples), maxFeatures_(maxFeatures), seed_(seed) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<int>& y) {
            assert(X.rows() == y.size());
            size_t n = X.rows();
            size_t totalFeatures = X.cols();
            size_t mFeatures = maxFeatures_ > 0
                ? maxFeatures_
                : static_cast<size_t>(std::max(1.0, std::sqrt(static_cast<double>(totalFeatures))));
            mFeatures = std::min(mFeatures, totalFeatures);

            std::mt19937 rng(seed_);
            trees_.clear();
            featureSubsets_.clear();
            trees_.resize(numTrees_);
            featureSubsets_.resize(numTrees_);

            std::vector<size_t> allFeatures(totalFeatures);
            std::iota(allFeatures.begin(), allFeatures.end(), 0);

            for (size_t t = 0; t < numTrees_; ++t) {
                // ブートストラップサンプリング (復元抽出)
                std::uniform_int_distribution<size_t> sampleDist(0, n - 1);
                std::vector<size_t> bootstrapIdx(n);
                for (size_t i = 0; i < n; ++i)
                    bootstrapIdx[i] = sampleDist(rng);

                // 特徴量サブサンプリング (非復元)
                std::vector<size_t> featPerm = allFeatures;
                std::shuffle(featPerm.begin(), featPerm.end(), rng);
                featureSubsets_[t].assign(featPerm.begin(), featPerm.begin() + mFeatures);
                std::sort(featureSubsets_[t].begin(), featureSubsets_[t].end());

                // 部分行列を構築
                Matrix<T> subX(n, mFeatures);
                std::vector<int> subY(n);
                for (size_t i = 0; i < n; ++i) {
                    size_t si = bootstrapIdx[i];
                    subY[i] = y[si];
                    for (size_t f = 0; f < mFeatures; ++f)
                        subX(i, f) = X(si, featureSubsets_[t][f]);
                }

                trees_[t] = DecisionTreeClassifier<T>(maxDepth_, minSamples_);
                trees_[t].fit(subX, subY);
            }
        }

        /// 単一サンプルの予測 (多数決)
        int predict(const Matrix<T>& X, size_t row) const {
            std::unordered_map<int, size_t> votes;
            for (size_t t = 0; t < numTrees_; ++t) {
                Matrix<T> subX(1, featureSubsets_[t].size());
                for (size_t f = 0; f < featureSubsets_[t].size(); ++f)
                    subX(0, f) = X(row, featureSubsets_[t][f]);
                votes[trees_[t].predict(subX, 0)]++;
            }
            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;
        }

        /// 木の数
        size_t numTrees() const { return numTrees_; }

    private:
        size_t numTrees_;
        size_t maxDepth_;
        size_t minSamples_;
        size_t maxFeatures_;
        unsigned seed_;
        std::vector<DecisionTreeClassifier<T>> trees_;
        std::vector<std::vector<size_t>> featureSubsets_;
    };

    // =====================================================================
    // RandomForestRegressor
    // =====================================================================

    /// ランダムフォレスト回帰器
    template<typename T>
    class RandomForestRegressor {
    public:
        /// @param numTrees      木の数
        /// @param maxDepth      各木の最大深さ (0 = 無制限)
        /// @param minSamples    分割に必要な最小サンプル数
        /// @param maxFeatures   各木で使う特徴量数 (0 = n/3)
        /// @param seed          乱数シード
        explicit RandomForestRegressor(size_t numTrees = 100,
                                        size_t maxDepth = 0,
                                        size_t minSamples = 2,
                                        size_t maxFeatures = 0,
                                        unsigned seed = 42)
            : numTrees_(numTrees), maxDepth_(maxDepth),
              minSamples_(minSamples), maxFeatures_(maxFeatures), seed_(seed) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<T>& y) {
            assert(X.rows() == y.size());
            size_t n = X.rows();
            size_t totalFeatures = X.cols();
            size_t mFeatures = maxFeatures_ > 0
                ? maxFeatures_
                : static_cast<size_t>(std::max(1.0, static_cast<double>(totalFeatures) / 3.0));
            mFeatures = std::min(mFeatures, totalFeatures);

            std::mt19937 rng(seed_);
            trees_.clear();
            featureSubsets_.clear();
            trees_.resize(numTrees_);
            featureSubsets_.resize(numTrees_);

            std::vector<size_t> allFeatures(totalFeatures);
            std::iota(allFeatures.begin(), allFeatures.end(), 0);

            for (size_t t = 0; t < numTrees_; ++t) {
                std::uniform_int_distribution<size_t> sampleDist(0, n - 1);
                std::vector<size_t> bootstrapIdx(n);
                for (size_t i = 0; i < n; ++i)
                    bootstrapIdx[i] = sampleDist(rng);

                std::vector<size_t> featPerm = allFeatures;
                std::shuffle(featPerm.begin(), featPerm.end(), rng);
                featureSubsets_[t].assign(featPerm.begin(), featPerm.begin() + mFeatures);
                std::sort(featureSubsets_[t].begin(), featureSubsets_[t].end());

                Matrix<T> subX(n, mFeatures);
                std::vector<T> subY(n);
                for (size_t i = 0; i < n; ++i) {
                    size_t si = bootstrapIdx[i];
                    subY[i] = y[si];
                    for (size_t f = 0; f < mFeatures; ++f)
                        subX(i, f) = X(si, featureSubsets_[t][f]);
                }

                trees_[t] = DecisionTreeRegressor<T>(maxDepth_, minSamples_);
                trees_[t].fit(subX, subY);
            }
        }

        /// 単一サンプルの予測 (平均)
        T predict(const Matrix<T>& X, size_t row) const {
            T sum = T(0);
            for (size_t t = 0; t < numTrees_; ++t) {
                Matrix<T> subX(1, featureSubsets_[t].size());
                for (size_t f = 0; f < featureSubsets_[t].size(); ++f)
                    subX(0, f) = X(row, featureSubsets_[t][f]);
                sum += trees_[t].predict(subX, 0);
            }
            return sum / static_cast<T>(numTrees_);
        }

        /// 全サンプルの予測
        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 numTrees() const { return numTrees_; }

    private:
        size_t numTrees_;
        size_t maxDepth_;
        size_t minSamples_;
        size_t maxFeatures_;
        unsigned seed_;
        std::vector<DecisionTreeRegressor<T>> trees_;
        std::vector<std::vector<size_t>> featureSubsets_;
    };

} // namespace calx::ml

#endif // CALX_ML_ENSEMBLE_RANDOMFOREST_HPP
