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

// Optimizer.hpp — オプティマイザ
//
// SGD (momentum), Adam, AdamW。
// Sequential::parameters() から取得した ParamGrad リストを更新する。

#ifndef CALX_ML_NN_OPTIMIZER_HPP
#define CALX_ML_NN_OPTIMIZER_HPP

#include <ml/nn/Layer.hpp>
#include <cmath>
#include <vector>

namespace calx::ml::nn {

    // =================================================================
    // OptimizerBase
    // =================================================================

    template<typename T>
    class OptimizerBase {
    public:
        virtual ~OptimizerBase() = default;

        /// パラメータを1ステップ更新
        virtual void step(std::vector<ParamGrad<T>>& params) = 0;

        /// 全勾配をゼロにリセット
        static void zeroGrad(std::vector<ParamGrad<T>>& params) {
            for (auto& pg : params) {
                auto& g = *pg.grad;
                for (size_t i = 0; i < g.rows(); ++i)
                    for (size_t j = 0; j < g.cols(); ++j)
                        g(i, j) = T{0};
            }
        }
    };

    // =================================================================
    // SGD (Stochastic Gradient Descent with Momentum)
    // =================================================================

    template<typename T>
    class SGD : public OptimizerBase<T> {
    public:
        /// @param lr        学習率
        /// @param momentum  モメンタム係数 (0 = なし)
        SGD(T lr = T{0.01}, T momentum = T{0})
            : lr_(lr), momentum_(momentum) {}

        void step(std::vector<ParamGrad<T>>& params) override {
            // 初回: velocity バッファを初期化
            if (momentum_ > T{0} && velocity_.empty()) {
                velocity_.resize(params.size());
                for (size_t p = 0; p < params.size(); ++p)
                    velocity_[p] = Matrix<T>(params[p].param->rows(),
                                              params[p].param->cols(), T{0});
            }

            for (size_t p = 0; p < params.size(); ++p) {
                auto& W = *params[p].param;
                auto& G = *params[p].grad;
                size_t rows = W.rows(), cols = W.cols();

                if (momentum_ > T{0}) {
                    auto& V = velocity_[p];
                    for (size_t i = 0; i < rows; ++i)
                        for (size_t j = 0; j < cols; ++j) {
                            V(i, j) = momentum_ * V(i, j) - lr_ * G(i, j);
                            W(i, j) += V(i, j);
                        }
                } else {
                    for (size_t i = 0; i < rows; ++i)
                        for (size_t j = 0; j < cols; ++j)
                            W(i, j) -= lr_ * G(i, j);
                }
            }
        }

    private:
        T lr_, momentum_;
        std::vector<Matrix<T>> velocity_;
    };

    // =================================================================
    // Adam
    // =================================================================

    template<typename T>
    class Adam : public OptimizerBase<T> {
    public:
        /// @param lr    学習率
        /// @param beta1 一次モーメント減衰率
        /// @param beta2 二次モーメント減衰率
        /// @param eps   数値安定化項
        Adam(T lr = T{0.001}, T beta1 = T{0.9}, T beta2 = T{0.999}, T eps = T{1e-8})
            : lr_(lr), beta1_(beta1), beta2_(beta2), eps_(eps) {}

        void step(std::vector<ParamGrad<T>>& params) override {
            if (m_.empty()) {
                m_.resize(params.size());
                v_.resize(params.size());
                for (size_t p = 0; p < params.size(); ++p) {
                    m_[p] = Matrix<T>(params[p].param->rows(), params[p].param->cols(), T{0});
                    v_[p] = Matrix<T>(params[p].param->rows(), params[p].param->cols(), T{0});
                }
            }

            ++t_;
            T bc1 = T{1} - std::pow(beta1_, static_cast<T>(t_));
            T bc2 = T{1} - std::pow(beta2_, static_cast<T>(t_));

            for (size_t p = 0; p < params.size(); ++p) {
                auto& W = *params[p].param;
                auto& G = *params[p].grad;
                auto& M = m_[p];
                auto& V = v_[p];
                size_t rows = W.rows(), cols = W.cols();

                for (size_t i = 0; i < rows; ++i)
                    for (size_t j = 0; j < cols; ++j) {
                        T g = G(i, j);
                        M(i, j) = beta1_ * M(i, j) + (T{1} - beta1_) * g;
                        V(i, j) = beta2_ * V(i, j) + (T{1} - beta2_) * g * g;
                        T mHat = M(i, j) / bc1;
                        T vHat = V(i, j) / bc2;
                        W(i, j) -= lr_ * mHat / (std::sqrt(vHat) + eps_);
                    }
            }
        }

    private:
        T lr_, beta1_, beta2_, eps_;
        int t_ = 0;
        std::vector<Matrix<T>> m_, v_;
    };

    // =================================================================
    // AdamW (Adam with decoupled weight decay)
    // =================================================================

    template<typename T>
    class AdamW : public OptimizerBase<T> {
    public:
        /// @param lr          学習率
        /// @param beta1       一次モーメント減衰率
        /// @param beta2       二次モーメント減衰率
        /// @param eps         数値安定化項
        /// @param weightDecay 重み減衰率
        AdamW(T lr = T{0.001}, T beta1 = T{0.9}, T beta2 = T{0.999},
              T eps = T{1e-8}, T weightDecay = T{0.01})
            : lr_(lr), beta1_(beta1), beta2_(beta2), eps_(eps), wd_(weightDecay) {}

        void step(std::vector<ParamGrad<T>>& params) override {
            if (m_.empty()) {
                m_.resize(params.size());
                v_.resize(params.size());
                for (size_t p = 0; p < params.size(); ++p) {
                    m_[p] = Matrix<T>(params[p].param->rows(), params[p].param->cols(), T{0});
                    v_[p] = Matrix<T>(params[p].param->rows(), params[p].param->cols(), T{0});
                }
            }

            ++t_;
            T bc1 = T{1} - std::pow(beta1_, static_cast<T>(t_));
            T bc2 = T{1} - std::pow(beta2_, static_cast<T>(t_));

            for (size_t p = 0; p < params.size(); ++p) {
                auto& W = *params[p].param;
                auto& G = *params[p].grad;
                auto& M = m_[p];
                auto& V = v_[p];
                size_t rows = W.rows(), cols = W.cols();

                for (size_t i = 0; i < rows; ++i)
                    for (size_t j = 0; j < cols; ++j) {
                        // Decoupled weight decay
                        W(i, j) -= lr_ * wd_ * W(i, j);

                        T g = G(i, j);
                        M(i, j) = beta1_ * M(i, j) + (T{1} - beta1_) * g;
                        V(i, j) = beta2_ * V(i, j) + (T{1} - beta2_) * g * g;
                        T mHat = M(i, j) / bc1;
                        T vHat = V(i, j) / bc2;
                        W(i, j) -= lr_ * mHat / (std::sqrt(vHat) + eps_);
                    }
            }
        }

    private:
        T lr_, beta1_, beta2_, eps_, wd_;
        int t_ = 0;
        std::vector<Matrix<T>> m_, v_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_OPTIMIZER_HPP
