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

// SVM.hpp — サポートベクターマシン
//
// SMO (Sequential Minimal Optimization) ソルバー。
// カーネル: 線形、RBF (ガウシアン)、多項式。
// 二値分類。多クラスは One-vs-One 方式。

#ifndef CALX_ML_SVM_SVM_HPP
#define CALX_ML_SVM_SVM_HPP

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

namespace calx::ml {

    /// カーネル種別
    enum class SvmKernel { Linear, RBF, Polynomial };

    namespace detail {

        /// カーネル関数の計算
        template<typename T>
        T kernelValue(SvmKernel type, const Matrix<T>& X, size_t i, size_t j,
                      size_t d, T gamma, T coef0, int degree)
        {
            T dot = T(0);
            for (size_t f = 0; f < d; ++f)
                dot += X(i, f) * X(j, f);

            switch (type) {
            case SvmKernel::Linear:
                return dot;
            case SvmKernel::RBF: {
                T dist = T(0);
                for (size_t f = 0; f < d; ++f) {
                    T diff = X(i, f) - X(j, f);
                    dist += diff * diff;
                }
                return std::exp(-gamma * dist);
            }
            case SvmKernel::Polynomial: {
                T val = gamma * dot + coef0;
                T result = T(1);
                for (int p = 0; p < degree; ++p) result *= val;
                return result;
            }
            }
            return dot;
        }

        /// 2 つのデータ行列間のカーネル値
        template<typename T>
        T kernelValueCross(SvmKernel type, const Matrix<T>& X1, size_t i,
                           const Matrix<T>& X2, size_t j, size_t d,
                           T gamma, T coef0, int degree)
        {
            T dot = T(0);
            for (size_t f = 0; f < d; ++f)
                dot += X1(i, f) * X2(j, f);

            switch (type) {
            case SvmKernel::Linear:
                return dot;
            case SvmKernel::RBF: {
                T dist = T(0);
                for (size_t f = 0; f < d; ++f) {
                    T diff = X1(i, f) - X2(j, f);
                    dist += diff * diff;
                }
                return std::exp(-gamma * dist);
            }
            case SvmKernel::Polynomial: {
                T val = gamma * dot + coef0;
                T result = T(1);
                for (int p = 0; p < degree; ++p) result *= val;
                return result;
            }
            }
            return dot;
        }

        // =================================================================
        // SMO 二値分類器
        // =================================================================
        template<typename T>
        class SmoBinaryClassifier {
        public:
            SmoBinaryClassifier() = default;

            void fit(const Matrix<T>& X, const std::vector<T>& y,
                     T C, SvmKernel kernel, T gamma, T coef0, int degree,
                     size_t maxIter, T tol)
            {
                size_t n = X.rows();
                size_t d = X.cols();
                alpha_.assign(n, T(0));
                b_ = T(0);

                // カーネル行列をキャッシュ (小規模データ向け)
                Matrix<T> K(n, n);
                for (size_t i = 0; i < n; ++i)
                    for (size_t j = i; j < n; ++j) {
                        T val = kernelValue(kernel, X, i, j, d, gamma, coef0, degree);
                        K(i, j) = val;
                        K(j, i) = val;
                    }

                // エラーキャッシュ
                std::vector<T> E(n);
                for (size_t i = 0; i < n; ++i) {
                    E[i] = -y[i];
                    for (size_t j = 0; j < n; ++j)
                        E[i] += alpha_[j] * y[j] * K(j, i);
                    E[i] += b_;  // = f(x_i) - y_i, initially -y_i
                }

                for (size_t iter = 0; iter < maxIter; ++iter) {
                    size_t numChanged = 0;
                    for (size_t i = 0; i < n; ++i) {
                        T ai = alpha_[i];
                        T yi = y[i];
                        T Ei = E[i];

                        // KKT 違反チェック
                        if ((yi * Ei < -tol && ai < C) ||
                            (yi * Ei > tol && ai > T(0)))
                        {
                            // j を選択 (|Ei - Ej| 最大)
                            size_t j = 0;
                            T maxDelta = T(0);
                            for (size_t k = 0; k < n; ++k) {
                                if (k == i) continue;
                                T delta = std::abs(Ei - E[k]);
                                if (delta > maxDelta) {
                                    maxDelta = delta;
                                    j = k;
                                }
                            }

                            T yj = y[j];
                            T Ej = E[j];
                            T oldAi = alpha_[i];
                            T oldAj = alpha_[j];

                            // 上下限
                            T L, H;
                            if (yi != yj) {
                                L = std::max(T(0), oldAj - oldAi);
                                H = std::min(C, C + oldAj - oldAi);
                            } else {
                                L = std::max(T(0), oldAi + oldAj - C);
                                H = std::min(C, oldAi + oldAj);
                            }
                            if (L >= H) continue;

                            T eta = T(2) * K(i, j) - K(i, i) - K(j, j);
                            if (eta >= T(0)) continue;

                            // alpha_j 更新
                            alpha_[j] = oldAj - yj * (Ei - Ej) / eta;
                            if (alpha_[j] > H) alpha_[j] = H;
                            if (alpha_[j] < L) alpha_[j] = L;

                            if (std::abs(alpha_[j] - oldAj) < T(1e-8))
                                continue;

                            // alpha_i 更新
                            alpha_[i] = oldAi + yi * yj * (oldAj - alpha_[j]);

                            // バイアス更新
                            T b1 = b_ - Ei
                                   - yi * (alpha_[i] - oldAi) * K(i, i)
                                   - yj * (alpha_[j] - oldAj) * K(i, j);
                            T b2 = b_ - Ej
                                   - yi * (alpha_[i] - oldAi) * K(i, j)
                                   - yj * (alpha_[j] - oldAj) * K(j, j);

                            if (alpha_[i] > T(0) && alpha_[i] < C)
                                b_ = b1;
                            else if (alpha_[j] > T(0) && alpha_[j] < C)
                                b_ = b2;
                            else
                                b_ = (b1 + b2) / T(2);

                            // エラーキャッシュ更新
                            for (size_t k = 0; k < n; ++k) {
                                E[k] += yi * (alpha_[i] - oldAi) * K(i, k)
                                      + yj * (alpha_[j] - oldAj) * K(j, k)
                                      + (b_ - (b_ - (b1 + b2) / T(2)));
                            }
                            // 再計算 (正確性のため)
                            for (size_t k = 0; k < n; ++k) {
                                E[k] = -y[k] + b_;
                                for (size_t m = 0; m < n; ++m)
                                    E[k] += alpha_[m] * y[m] * K(m, k);
                            }

                            ++numChanged;
                        }
                    }
                    if (numChanged == 0) break;
                }

                // サポートベクトルを記録
                supportVectors_.clear();
                supportAlpha_.clear();
                supportY_.clear();
                size_t d2 = X.cols();
                for (size_t i = 0; i < n; ++i) {
                    if (alpha_[i] > T(1e-8)) {
                        supportY_.push_back(y[i]);
                        supportAlpha_.push_back(alpha_[i]);
                        // 行をコピー
                        std::vector<T> row(d2);
                        for (size_t f = 0; f < d2; ++f) row[f] = X(i, f);
                        supportVectors_.push_back(std::move(row));
                    }
                }

                kernel_ = kernel;
                gamma_ = gamma;
                coef0_ = coef0;
                degree_ = degree;
                dim_ = d2;
            }

            T decisionFunction(const Matrix<T>& X, size_t row) const {
                T result = b_;
                for (size_t s = 0; s < supportVectors_.size(); ++s) {
                    T kval = kernelValueSV(X, row, s);
                    result += supportAlpha_[s] * supportY_[s] * kval;
                }
                return result;
            }

            int predict(const Matrix<T>& X, size_t row) const {
                return decisionFunction(X, row) >= T(0) ? 1 : -1;
            }

            size_t numSupportVectors() const { return supportVectors_.size(); }

        private:
            std::vector<T> alpha_;
            T b_ = T(0);

            std::vector<std::vector<T>> supportVectors_;
            std::vector<T> supportAlpha_;
            std::vector<T> supportY_;

            SvmKernel kernel_ = SvmKernel::Linear;
            T gamma_ = T(1);
            T coef0_ = T(0);
            int degree_ = 3;
            size_t dim_ = 0;

            T kernelValueSV(const Matrix<T>& X, size_t row, size_t svIdx) const {
                const auto& sv = supportVectors_[svIdx];
                T dot = T(0);
                for (size_t f = 0; f < dim_; ++f)
                    dot += X(row, f) * sv[f];

                switch (kernel_) {
                case SvmKernel::Linear:
                    return dot;
                case SvmKernel::RBF: {
                    T dist = T(0);
                    for (size_t f = 0; f < dim_; ++f) {
                        T diff = X(row, f) - sv[f];
                        dist += diff * diff;
                    }
                    return std::exp(-gamma_ * dist);
                }
                case SvmKernel::Polynomial: {
                    T val = gamma_ * dot + coef0_;
                    T result = T(1);
                    for (int p = 0; p < degree_; ++p) result *= val;
                    return result;
                }
                }
                return dot;
            }
        };

    } // namespace detail

    // =====================================================================
    // SVM (多クラス One-vs-One)
    // =====================================================================

    /// サポートベクターマシン分類器
    template<typename T>
    class SVM {
    public:
        /// @param C        正則化パラメータ (ソフトマージン)
        /// @param kernel   カーネル種別
        /// @param gamma    RBF/Polynomial の gamma (0 = 1/d)
        /// @param coef0    Polynomial の定数項
        /// @param degree   Polynomial の次数
        /// @param maxIter  SMO 最大反復数
        /// @param tol      KKT 違反の許容誤差
        explicit SVM(T C = T(1),
                     SvmKernel kernel = SvmKernel::RBF,
                     T gamma = T(0),
                     T coef0 = T(0),
                     int degree = 3,
                     size_t maxIter = 1000,
                     T tol = T(1e-3))
            : C_(C), kernel_(kernel), gamma_(gamma), coef0_(coef0),
              degree_(degree), maxIter_(maxIter), tol_(tol) {}

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

            T gamma = gamma_ > T(0) ? gamma_ : T(1) / static_cast<T>(d);

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

            if (K == 2) {
                // 二値: 直接 SMO
                size_t n = X.rows();
                std::vector<T> ySigned(n);
                for (size_t i = 0; i < n; ++i)
                    ySigned[i] = (y[i] == classes_[1]) ? T(1) : T(-1);

                binaryClassifier_ = detail::SmoBinaryClassifier<T>();
                binaryClassifier_.fit(X, ySigned, C_, kernel_, gamma, coef0_,
                                       degree_, maxIter_, tol_);
                isBinary_ = true;
            } else {
                // 多クラス: One-vs-One
                isBinary_ = false;
                ovoClassifiers_.clear();
                ovoPairs_.clear();

                for (size_t a = 0; a < K; ++a) {
                    for (size_t b = a + 1; b < K; ++b) {
                        // クラス a, b のサンプルのみ抽出
                        std::vector<size_t> indices;
                        for (size_t i = 0; i < X.rows(); ++i) {
                            if (y[i] == classes_[a] || y[i] == classes_[b])
                                indices.push_back(i);
                        }

                        Matrix<T> subX(indices.size(), d);
                        std::vector<T> subY(indices.size());
                        for (size_t i = 0; i < indices.size(); ++i) {
                            for (size_t f = 0; f < d; ++f)
                                subX(i, f) = X(indices[i], f);
                            subY[i] = (y[indices[i]] == classes_[b]) ? T(1) : T(-1);
                        }

                        detail::SmoBinaryClassifier<T> clf;
                        clf.fit(subX, subY, C_, kernel_, gamma, coef0_,
                                degree_, maxIter_, tol_);
                        ovoClassifiers_.push_back(std::move(clf));
                        ovoPairs_.push_back({ static_cast<int>(a), static_cast<int>(b) });
                    }
                }
            }
        }

        /// 単一サンプルの予測
        int predict(const Matrix<T>& X, size_t row) const {
            if (isBinary_) {
                int sign = binaryClassifier_.predict(X, row);
                return sign >= 0 ? classes_[1] : classes_[0];
            } else {
                // 投票
                std::vector<int> votes(classes_.size(), 0);
                for (size_t p = 0; p < ovoClassifiers_.size(); ++p) {
                    int sign = ovoClassifiers_[p].predict(X, row);
                    if (sign >= 0)
                        votes[ovoPairs_[p].second]++;
                    else
                        votes[ovoPairs_[p].first]++;
                }
                auto maxIt = std::max_element(votes.begin(), votes.end());
                return classes_[std::distance(votes.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 C_;
        SvmKernel kernel_;
        T gamma_;
        T coef0_;
        int degree_;
        size_t maxIter_;
        T tol_;

        std::vector<int> classes_;
        bool isBinary_ = true;

        // 二値用
        detail::SmoBinaryClassifier<T> binaryClassifier_;

        // 多クラス用 (One-vs-One)
        struct ClassPair { int first, second; };
        std::vector<detail::SmoBinaryClassifier<T>> ovoClassifiers_;
        std::vector<ClassPair> ovoPairs_;
    };

} // namespace calx::ml

#endif // CALX_ML_SVM_SVM_HPP
