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

// GradientBoosting.hpp — 勾配ブースティング決定木 (GBDT)
//
// 回帰: MSE 損失 (残差フィッティング)
// 分類: 二値 log-loss → 多クラス One-vs-Rest
// 各ステージで浅い回帰木を残差/勾配にフィットさせる。

#ifndef CALX_ML_ENSEMBLE_GRADIENTBOOSTING_HPP
#define CALX_ML_ENSEMBLE_GRADIENTBOOSTING_HPP

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

namespace calx::ml {

    // =====================================================================
    // GradientBoostingRegressor
    // =====================================================================

    /// GBDT 回帰器 (MSE 損失)
    template<typename T>
    class GradientBoostingRegressor {
    public:
        /// @param numTrees       ブースティングラウンド数
        /// @param learningRate   学習率 (shrinkage)
        /// @param maxDepth       各木の最大深さ
        /// @param minSamples     分割に必要な最小サンプル数
        explicit GradientBoostingRegressor(size_t numTrees = 100,
                                            T learningRate = T(0.1),
                                            size_t maxDepth = 3,
                                            size_t minSamples = 2)
            : numTrees_(numTrees), learningRate_(learningRate),
              maxDepth_(maxDepth), minSamples_(minSamples) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<T>& y) {
            assert(X.rows() == y.size());
            size_t n = X.rows();

            // 初期予測 = 平均
            T mean = T(0);
            for (auto v : y) mean += v;
            mean /= static_cast<T>(n);
            initialPrediction_ = mean;

            trees_.clear();
            trees_.reserve(numTrees_);

            // 現在の予測値
            std::vector<T> predictions(n, mean);
            std::vector<T> residuals(n);

            for (size_t t = 0; t < numTrees_; ++t) {
                // 負の勾配 = 残差 (MSE 損失)
                for (size_t i = 0; i < n; ++i)
                    residuals[i] = y[i] - predictions[i];

                // 残差に回帰木をフィット
                DecisionTreeRegressor<T> tree(maxDepth_, minSamples_);
                tree.fit(X, residuals);

                // 予測を更新
                auto treePred = tree.predict(X);
                for (size_t i = 0; i < n; ++i)
                    predictions[i] += learningRate_ * treePred[i];

                trees_.push_back(std::move(tree));
            }
        }

        /// 単一サンプルの予測
        T predict(const Matrix<T>& X, size_t row) const {
            T pred = initialPrediction_;
            for (auto& tree : trees_)
                pred += learningRate_ * tree.predict(X, row);
            return pred;
        }

        /// 全サンプルの予測
        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 trees_.size(); }

    private:
        size_t numTrees_;
        T learningRate_;
        size_t maxDepth_;
        size_t minSamples_;
        T initialPrediction_ = T(0);
        std::vector<DecisionTreeRegressor<T>> trees_;
    };

    // =====================================================================
    // GradientBoostingClassifier
    // =====================================================================

    /// GBDT 分類器 (多クラス: One-vs-Rest log-loss)
    template<typename T>
    class GradientBoostingClassifier {
    public:
        /// @param numTrees       ブースティングラウンド数
        /// @param learningRate   学習率 (shrinkage)
        /// @param maxDepth       各木の最大深さ
        /// @param minSamples     分割に必要な最小サンプル数
        explicit GradientBoostingClassifier(size_t numTrees = 100,
                                             T learningRate = T(0.1),
                                             size_t maxDepth = 3,
                                             size_t minSamples = 2)
            : numTrees_(numTrees), learningRate_(learningRate),
              maxDepth_(maxDepth), minSamples_(minSamples) {}

        /// 学習
        void fit(const Matrix<T>& X, const std::vector<int>& y) {
            assert(X.rows() == y.size());
            size_t n = X.rows();

            // クラスラベル収集
            std::set<int> classSet(y.begin(), y.end());
            classes_.assign(classSet.begin(), classSet.end());
            size_t K = classes_.size();

            if (K == 2) {
                // 二値分類: 単一の log-odds モデル
                fitBinary(X, y, n);
            } else {
                // 多クラス: One-vs-Rest
                fitMulticlass(X, y, n, K);
            }
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            if (classes_.size() == 2) {
                T logOdds = initialLogOdds_;
                for (auto& tree : binaryTrees_)
                    logOdds += learningRate_ * tree.predict(X, row);
                T prob = sigmoid(logOdds);
                return prob >= T(0.5) ? classes_[1] : classes_[0];
            } else {
                // 多クラス: 各クラスのスコアを計算
                size_t K = classes_.size();
                std::vector<T> scores(K);
                for (size_t k = 0; k < K; ++k) {
                    scores[k] = initialScores_[k];
                    for (auto& tree : multiTrees_[k])
                        scores[k] += learningRate_ * tree.predict(X, row);
                }
                auto maxIt = std::max_element(scores.begin(), scores.end());
                return classes_[std::distance(scores.begin(), maxIt)];
            }
        }

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

        const std::vector<int>& classes() const { return classes_; }
        size_t numTrees() const { return numTrees_; }

    private:
        size_t numTrees_;
        T learningRate_;
        size_t maxDepth_;
        size_t minSamples_;

        std::vector<int> classes_;

        // 二値分類用
        T initialLogOdds_ = T(0);
        std::vector<DecisionTreeRegressor<T>> binaryTrees_;

        // 多クラス用
        std::vector<T> initialScores_;
        std::vector<std::vector<DecisionTreeRegressor<T>>> multiTrees_;

        static T sigmoid(T x) {
            return T(1) / (T(1) + std::exp(-x));
        }

        static T clamp(T x, T lo, T hi) {
            return x < lo ? lo : (x > hi ? hi : x);
        }

        void fitBinary(const Matrix<T>& X, const std::vector<int>& y, size_t n) {
            // ターゲット: classes_[1] → 1, classes_[0] → 0
            int posClass = classes_[1];
            std::vector<T> target(n);
            T posCount = T(0);
            for (size_t i = 0; i < n; ++i) {
                target[i] = (y[i] == posClass) ? T(1) : T(0);
                posCount += target[i];
            }

            // 初期 log-odds
            T p = clamp(posCount / static_cast<T>(n), T(1e-7), T(1 - 1e-7));
            initialLogOdds_ = std::log(p / (T(1) - p));

            binaryTrees_.clear();
            binaryTrees_.reserve(numTrees_);

            std::vector<T> predictions(n, initialLogOdds_);
            std::vector<T> residuals(n);

            for (size_t t = 0; t < numTrees_; ++t) {
                // 負の勾配 = target - sigmoid(prediction)
                for (size_t i = 0; i < n; ++i) {
                    T prob = sigmoid(predictions[i]);
                    residuals[i] = target[i] - prob;
                }

                DecisionTreeRegressor<T> tree(maxDepth_, minSamples_);
                tree.fit(X, residuals);

                auto treePred = tree.predict(X);
                for (size_t i = 0; i < n; ++i)
                    predictions[i] += learningRate_ * treePred[i];

                binaryTrees_.push_back(std::move(tree));
            }
        }

        void fitMulticlass(const Matrix<T>& X, const std::vector<int>& y,
                           size_t n, size_t K)
        {
            // One-vs-Rest: 各クラスに対して二値回帰木を学習
            initialScores_.resize(K);
            multiTrees_.resize(K);

            for (size_t k = 0; k < K; ++k) {
                std::vector<T> target(n);
                T posCount = T(0);
                for (size_t i = 0; i < n; ++i) {
                    target[i] = (y[i] == classes_[k]) ? T(1) : T(0);
                    posCount += target[i];
                }

                T p = clamp(posCount / static_cast<T>(n), T(1e-7), T(1 - 1e-7));
                initialScores_[k] = std::log(p / (T(1) - p));

                multiTrees_[k].reserve(numTrees_);
                std::vector<T> predictions(n, initialScores_[k]);
                std::vector<T> residuals(n);

                for (size_t t = 0; t < numTrees_; ++t) {
                    for (size_t i = 0; i < n; ++i) {
                        T prob = sigmoid(predictions[i]);
                        residuals[i] = target[i] - prob;
                    }

                    DecisionTreeRegressor<T> tree(maxDepth_, minSamples_);
                    tree.fit(X, residuals);

                    auto treePred = tree.predict(X);
                    for (size_t i = 0; i < n; ++i)
                        predictions[i] += learningRate_ * treePred[i];

                    multiTrees_[k].push_back(std::move(tree));
                }
            }
        }
    };

} // namespace calx::ml

#endif // CALX_ML_ENSEMBLE_GRADIENTBOOSTING_HPP
