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

// QLearning.hpp — Q-Learning (テーブル方式)
//
// Q(s,a) ← Q(s,a) + α[r + γ max_a' Q(s',a') - Q(s,a)]

#ifndef CALX_ML_RL_QLEARNING_HPP
#define CALX_ML_RL_QLEARNING_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <random>
#include <vector>

namespace calx::ml {

    /// Q-Learning (テーブル方式)
    template<typename T>
    class QLearning {
    public:
        /// @param numStates   状態数
        /// @param numActions  行動数
        /// @param alpha       学習率
        /// @param gamma       割引率
        /// @param epsilon     ε-greedy の探索率
        /// @param seed        乱数シード
        QLearning(size_t numStates, size_t numActions,
                  T alpha = T{0.1}, T gamma = T{0.99}, T epsilon = T{0.1},
                  unsigned seed = 42)
            : nStates_(numStates), nActions_(numActions)
            , alpha_(alpha), gamma_(gamma), epsilon_(epsilon), rng_(seed)
        {
            Q_ = Matrix<T>(numStates, numActions, T{0});
        }

        /// ε-greedy で行動選択
        size_t selectAction(size_t state) {
            std::uniform_real_distribution<T> dist(T{0}, T{1});
            if (dist(rng_) < epsilon_) {
                std::uniform_int_distribution<size_t> actDist(0, nActions_ - 1);
                return actDist(rng_);
            }
            return greedyAction(state);
        }

        /// Greedy 行動選択 (推論用)
        size_t greedyAction(size_t state) const {
            size_t bestA = 0;
            T bestQ = Q_(state, 0);
            for (size_t a = 1; a < nActions_; ++a)
                if (Q_(state, a) > bestQ) { bestQ = Q_(state, a); bestA = a; }
            return bestA;
        }

        /// Q 値更新
        void update(size_t state, size_t action, T reward, size_t nextState, bool done) {
            T maxNextQ = T{0};
            if (!done) {
                maxNextQ = Q_(nextState, 0);
                for (size_t a = 1; a < nActions_; ++a)
                    if (Q_(nextState, a) > maxNextQ) maxNextQ = Q_(nextState, a);
            }
            T target = reward + gamma_ * maxNextQ;
            Q_(state, action) += alpha_ * (target - Q_(state, action));
        }

        /// Q テーブル
        const Matrix<T>& qTable() const { return Q_; }
        T qValue(size_t state, size_t action) const { return Q_(state, action); }

        void setEpsilon(T eps) { epsilon_ = eps; }

    private:
        size_t nStates_, nActions_;
        T alpha_, gamma_, epsilon_;
        std::mt19937 rng_;
        Matrix<T> Q_;
    };

} // namespace calx::ml

#endif // CALX_ML_RL_QLEARNING_HPP
