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

// LogisticRegression.hpp — ロジスティック回帰
//
// 二値分類: sigmoid + binary cross-entropy + 勾配降下法
// 多クラス分類: softmax + cross-entropy (One-vs-Rest ではなく真の多クラス)
// L2 正則化対応。

#ifndef CALX_ML_LINEAR_LOGISTICREGRESSION_HPP
#define CALX_ML_LINEAR_LOGISTICREGRESSION_HPP

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

namespace calx::ml {

    /// ロジスティック回帰 (二値 + 多クラス自動切替)
    template<typename T>
    class LogisticRegression {
    public:
        /// @param learningRate  学習率
        /// @param maxIter       最大反復回数
        /// @param lambda        L2 正則化係数 (0 = なし)
        /// @param tol           収束判定閾値 (勾配ノルム)
        explicit LogisticRegression(T learningRate = T(0.01),
                                     size_t maxIter = 1000,
                                     T lambda = T(0),
                                     T tol = T(1e-6))
            : learningRate_(learningRate), maxIter_(maxIter),
              lambda_(lambda), tol_(tol) {}

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

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

            if (classes_.size() == 2) {
                fitBinary(X, y, n, d);
            } else {
                fitMulticlass(X, y, n, d);
            }
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            if (classes_.size() == 2) {
                T z = bias_;
                for (size_t j = 0; j < weights_.size(); ++j)
                    z += weights_[j] * X(row, j);
                return sigmoid(z) >= T(0.5) ? classes_[1] : classes_[0];
            } else {
                auto probs = predictProbability(X, row);
                auto maxIt = std::max_element(probs.begin(), probs.end());
                return classes_[std::distance(probs.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;
        }

        /// クラス確率 (二値: {P(class0), P(class1)}, 多クラス: softmax)
        std::vector<T> predictProbability(const Matrix<T>& X, size_t row) const {
            if (classes_.size() == 2) {
                T z = bias_;
                for (size_t j = 0; j < weights_.size(); ++j)
                    z += weights_[j] * X(row, j);
                T p1 = sigmoid(z);
                return { T(1) - p1, p1 };
            } else {
                size_t K = classes_.size();
                size_t d = weights_.size() / K;
                std::vector<T> logits(K);
                for (size_t k = 0; k < K; ++k) {
                    logits[k] = biases_[k];
                    for (size_t j = 0; j < d; ++j)
                        logits[k] += weights_[k * d + j] * X(row, j);
                }
                return softmax(logits);
            }
        }

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

        /// 収束時の反復回数
        size_t iterations() const { return iterations_; }

    private:
        T learningRate_;
        size_t maxIter_;
        T lambda_;
        T tol_;

        std::vector<int> classes_;

        // 二値分類用
        std::vector<T> weights_;
        T bias_ = T(0);

        // 多クラス用
        std::vector<T> biases_;

        size_t iterations_ = 0;

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

        static std::vector<T> softmax(const std::vector<T>& logits) {
            T maxVal = *std::max_element(logits.begin(), logits.end());
            std::vector<T> result(logits.size());
            T sum = T(0);
            for (size_t i = 0; i < logits.size(); ++i) {
                result[i] = std::exp(logits[i] - maxVal);
                sum += result[i];
            }
            for (auto& v : result) v /= sum;
            return result;
        }

        void fitBinary(const Matrix<T>& X, const std::vector<int>& y,
                        size_t n, size_t d)
        {
            int posClass = classes_[1];
            weights_.assign(d, T(0));
            bias_ = T(0);

            for (size_t iter = 0; iter < maxIter_; ++iter) {
                std::vector<T> gradW(d, T(0));
                T gradB = T(0);

                for (size_t i = 0; i < n; ++i) {
                    T z = bias_;
                    for (size_t j = 0; j < d; ++j)
                        z += weights_[j] * X(i, j);
                    T p = sigmoid(z);
                    T target = (y[i] == posClass) ? T(1) : T(0);
                    T err = p - target;
                    for (size_t j = 0; j < d; ++j)
                        gradW[j] += err * X(i, j);
                    gradB += err;
                }

                // 平均勾配 + L2 正則化
                T invN = T(1) / static_cast<T>(n);
                T gradNorm = T(0);
                for (size_t j = 0; j < d; ++j) {
                    gradW[j] = gradW[j] * invN + lambda_ * weights_[j];
                    gradNorm += gradW[j] * gradW[j];
                }
                gradB *= invN;
                gradNorm += gradB * gradB;

                if (std::sqrt(gradNorm) < tol_) {
                    iterations_ = iter + 1;
                    return;
                }

                for (size_t j = 0; j < d; ++j)
                    weights_[j] -= learningRate_ * gradW[j];
                bias_ -= learningRate_ * gradB;
            }
            iterations_ = maxIter_;
        }

        void fitMulticlass(const Matrix<T>& X, const std::vector<int>& y,
                            size_t n, size_t d)
        {
            size_t K = classes_.size();
            weights_.assign(K * d, T(0));
            biases_.assign(K, T(0));

            // ラベル → クラスインデックス
            std::vector<size_t> yIdx(n);
            for (size_t i = 0; i < n; ++i) {
                auto it = std::find(classes_.begin(), classes_.end(), y[i]);
                yIdx[i] = static_cast<size_t>(std::distance(classes_.begin(), it));
            }

            for (size_t iter = 0; iter < maxIter_; ++iter) {
                std::vector<T> gradW(K * d, T(0));
                std::vector<T> gradB(K, T(0));

                for (size_t i = 0; i < n; ++i) {
                    // logits
                    std::vector<T> logits(K);
                    for (size_t k = 0; k < K; ++k) {
                        logits[k] = biases_[k];
                        for (size_t j = 0; j < d; ++j)
                            logits[k] += weights_[k * d + j] * X(i, j);
                    }
                    auto probs = softmax(logits);

                    for (size_t k = 0; k < K; ++k) {
                        T err = probs[k] - ((yIdx[i] == k) ? T(1) : T(0));
                        for (size_t j = 0; j < d; ++j)
                            gradW[k * d + j] += err * X(i, j);
                        gradB[k] += err;
                    }
                }

                T invN = T(1) / static_cast<T>(n);
                T gradNorm = T(0);
                for (size_t k = 0; k < K; ++k) {
                    for (size_t j = 0; j < d; ++j) {
                        gradW[k * d + j] = gradW[k * d + j] * invN
                                           + lambda_ * weights_[k * d + j];
                        gradNorm += gradW[k * d + j] * gradW[k * d + j];
                    }
                    gradB[k] *= invN;
                    gradNorm += gradB[k] * gradB[k];
                }

                if (std::sqrt(gradNorm) < tol_) {
                    iterations_ = iter + 1;
                    return;
                }

                for (size_t idx = 0; idx < K * d; ++idx)
                    weights_[idx] -= learningRate_ * gradW[idx];
                for (size_t k = 0; k < K; ++k)
                    biases_[k] -= learningRate_ * gradB[k];
            }
            iterations_ = maxIter_;
        }
    };

} // namespace calx::ml

#endif // CALX_ML_LINEAR_LOGISTICREGRESSION_HPP
