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

// NaiveBayes.hpp — ナイーブベイズ分類器
//
// GaussianNB: 連続特徴量 (正規分布仮定)
// MultinomialNB: カウント/頻度特徴量 (多項分布仮定、Laplace 平滑化)

#ifndef CALX_ML_BAYES_NAIVEBAYES_HPP
#define CALX_ML_BAYES_NAIVEBAYES_HPP

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

namespace calx::ml {

    // =====================================================================
    // GaussianNB
    // =====================================================================

    /// ガウシアンナイーブベイズ分類器
    template<typename T>
    class GaussianNB {
    public:
        /// 学習
        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());
            size_t K = classes_.size();

            // 各クラスの平均・分散・事前確率を計算
            means_.assign(K * d, T(0));
            vars_.assign(K * d, T(0));
            logPriors_.resize(K);

            std::vector<size_t> counts(K, 0);

            for (size_t i = 0; i < n; ++i) {
                size_t k = classIndex(y[i]);
                counts[k]++;
                for (size_t j = 0; j < d; ++j)
                    means_[k * d + j] += X(i, j);
            }

            for (size_t k = 0; k < K; ++k) {
                T invCount = T(1) / static_cast<T>(counts[k]);
                for (size_t j = 0; j < d; ++j)
                    means_[k * d + j] *= invCount;
                logPriors_[k] = std::log(static_cast<T>(counts[k]) / static_cast<T>(n));
            }

            // 分散
            for (size_t i = 0; i < n; ++i) {
                size_t k = classIndex(y[i]);
                for (size_t j = 0; j < d; ++j) {
                    T diff = X(i, j) - means_[k * d + j];
                    vars_[k * d + j] += diff * diff;
                }
            }
            for (size_t k = 0; k < K; ++k) {
                T invCount = T(1) / static_cast<T>(counts[k]);
                for (size_t j = 0; j < d; ++j) {
                    vars_[k * d + j] *= invCount;
                    // ゼロ分散を防止
                    if (vars_[k * d + j] < T(1e-9))
                        vars_[k * d + j] = T(1e-9);
                }
            }

            dim_ = d;
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            auto logProbs = logPosterior(X, row);
            auto maxIt = std::max_element(logProbs.begin(), logProbs.end());
            return classes_[std::distance(logProbs.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_; }

    private:
        std::vector<int> classes_;
        std::vector<T> means_;     // K * d
        std::vector<T> vars_;      // K * d
        std::vector<T> logPriors_; // K
        size_t dim_ = 0;

        size_t classIndex(int label) const {
            auto it = std::find(classes_.begin(), classes_.end(), label);
            return static_cast<size_t>(std::distance(classes_.begin(), it));
        }

        /// 各クラスの対数事後確率 (非正規化)
        std::vector<T> logPosterior(const Matrix<T>& X, size_t row) const {
            size_t K = classes_.size();
            std::vector<T> logProbs(K);
            constexpr T log2pi = T(1.8378770664093453);  // log(2*pi)

            for (size_t k = 0; k < K; ++k) {
                logProbs[k] = logPriors_[k];
                for (size_t j = 0; j < dim_; ++j) {
                    T diff = X(row, j) - means_[k * dim_ + j];
                    T var = vars_[k * dim_ + j];
                    // log N(x | mu, sigma^2) = -0.5*(log(2pi) + log(var) + (x-mu)^2/var)
                    logProbs[k] -= T(0.5) * (log2pi + std::log(var) + diff * diff / var);
                }
            }
            return logProbs;
        }
    };

    // =====================================================================
    // MultinomialNB
    // =====================================================================

    /// 多項ナイーブベイズ分類器 (Laplace 平滑化付き)
    template<typename T>
    class MultinomialNB {
    public:
        /// @param alpha  Laplace 平滑化パラメータ (1.0 = Laplace, 0 = なし)
        explicit MultinomialNB(T alpha = T(1)) : alpha_(alpha) {}

        /// 学習
        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());
            size_t K = classes_.size();

            // 各クラスの特徴量合計を計算
            featureLogProb_.assign(K * d, T(0));
            logPriors_.resize(K);
            std::vector<size_t> counts(K, 0);

            for (size_t i = 0; i < n; ++i) {
                size_t k = classIndex(y[i]);
                counts[k]++;
                for (size_t j = 0; j < d; ++j)
                    featureLogProb_[k * d + j] += X(i, j);
            }

            // 対数条件付き確率 (Laplace 平滑化)
            for (size_t k = 0; k < K; ++k) {
                T totalCount = T(0);
                for (size_t j = 0; j < d; ++j)
                    totalCount += featureLogProb_[k * d + j];
                T denom = totalCount + alpha_ * static_cast<T>(d);
                for (size_t j = 0; j < d; ++j)
                    featureLogProb_[k * d + j] =
                        std::log((featureLogProb_[k * d + j] + alpha_) / denom);

                logPriors_[k] = std::log(static_cast<T>(counts[k]) / static_cast<T>(n));
            }

            dim_ = d;
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            size_t K = classes_.size();
            std::vector<T> logProbs(K);
            for (size_t k = 0; k < K; ++k) {
                logProbs[k] = logPriors_[k];
                for (size_t j = 0; j < dim_; ++j)
                    logProbs[k] += X(row, j) * featureLogProb_[k * dim_ + j];
            }
            auto maxIt = std::max_element(logProbs.begin(), logProbs.end());
            return classes_[std::distance(logProbs.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_; }

    private:
        T alpha_;
        std::vector<int> classes_;
        std::vector<T> featureLogProb_; // K * d
        std::vector<T> logPriors_;      // K
        size_t dim_ = 0;

        size_t classIndex(int label) const {
            auto it = std::find(classes_.begin(), classes_.end(), label);
            return static_cast<size_t>(std::distance(classes_.begin(), it));
        }
    };

} // namespace calx::ml

#endif // CALX_ML_BAYES_NAIVEBAYES_HPP
