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

// PolicyGradient.hpp — REINFORCE アルゴリズム
//
// 方策勾配法: ∇J(θ) = E[Σ_t ∇log π(a_t|s_t) * G_t]

#ifndef CALX_ML_RL_POLICYGRADIENT_HPP
#define CALX_ML_RL_POLICYGRADIENT_HPP

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

namespace calx::ml {

    /// REINFORCE (Policy Gradient)
    template<typename T>
    class REINFORCE {
    public:
        /// @param stateDim    状態次元
        /// @param numActions  行動数
        /// @param hiddenDim   隠れ層次元
        /// @param lr          学習率
        /// @param gamma       割引率
        /// @param seed        乱数シード
        REINFORCE(size_t stateDim, size_t numActions, size_t hiddenDim = 64,
                  T lr = T{0.001}, T gamma = T{0.99}, unsigned seed = 42)
            : stateDim_(stateDim), nActions_(numActions), gamma_(gamma), rng_(seed)
            , net1_(stateDim, hiddenDim, true, seed)
            , net2_(hiddenDim, numActions, true, seed + 1)
            , relu_(), softmax_()
            , optimizer_(lr)
        {}

        /// 方策に従い行動選択 (softmax サンプリング)
        size_t selectAction(const std::vector<T>& state) {
            Matrix<T> input(1, stateDim_);
            for (size_t d = 0; d < stateDim_; ++d) input(0, d) = state[d];

            auto probs = forwardPolicy(input);
            lastProbs_.resize(nActions_);
            for (size_t a = 0; a < nActions_; ++a) lastProbs_[a] = probs(0, a);

            // カテゴリカル分布からサンプリング
            std::uniform_real_distribution<T> dist(T{0}, T{1});
            T r = dist(rng_);
            T cumProb = T{0};
            for (size_t a = 0; a < nActions_; ++a) {
                cumProb += lastProbs_[a];
                if (r < cumProb) return a;
            }
            return nActions_ - 1;
        }

        /// エピソードの経験を記録
        void storeTransition(const std::vector<T>& state, size_t action, T reward) {
            states_.push_back(state);
            actions_.push_back(action);
            rewards_.push_back(reward);
        }

        /// エピソード終了時に方策を更新
        /// @return 平均損失
        T updatePolicy() {
            if (rewards_.empty()) return T{0};

            size_t n = rewards_.size();

            // 割引報酬の計算
            std::vector<T> returns(n);
            T G = T{0};
            for (int t = static_cast<int>(n) - 1; t >= 0; --t) {
                G = rewards_[t] + gamma_ * G;
                returns[t] = G;
            }

            // 報酬の正規化
            T mean = T{0}, var = T{0};
            for (auto r : returns) mean += r;
            mean /= static_cast<T>(n);
            for (auto r : returns) var += (r - mean) * (r - mean);
            var /= static_cast<T>(n);
            T invStd = (var > T{1e-8}) ? T{1} / std::sqrt(var) : T{1};
            for (auto& r : returns) r = (r - mean) * invStd;

            // 方策勾配
            auto params = parameters();
            T totalLoss = T{0};

            for (size_t t = 0; t < n; ++t) {
                Matrix<T> input(1, stateDim_);
                for (size_t d = 0; d < stateDim_; ++d) input(0, d) = states_[t][d];

                auto probs = forwardPolicy(input);

                // -log π(a|s) * G_t の勾配
                // Softmax の勾配: grad_j = π_j - 1(j==a) を returns で重み付け
                Matrix<T> grad(1, nActions_);
                for (size_t a = 0; a < nActions_; ++a) {
                    T indicator = (a == actions_[t]) ? T{1} : T{0};
                    grad(0, a) = (probs(0, a) - indicator) * returns[t];
                }

                nn::OptimizerBase<T>::zeroGrad(params);
                softmax_.backward(grad);
                auto d2 = net2_.backward(grad);
                net1_.backward(relu_.backward(d2));
                optimizer_.step(params);

                T prob = probs(0, actions_[t]);
                if (prob > T{1e-10})
                    totalLoss -= std::log(prob) * returns[t];
            }

            // バッファクリア
            states_.clear();
            actions_.clear();
            rewards_.clear();

            return totalLoss / static_cast<T>(n);
        }

        std::vector<nn::ParamGrad<T>> parameters() {
            auto p = net1_.parameters();
            auto p2 = net2_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            return p;
        }

    private:
        Matrix<T> forwardPolicy(const Matrix<T>& input) {
            auto h = relu_.forward(net1_.forward(input));
            return softmax_.forward(net2_.forward(h));
        }

        size_t stateDim_, nActions_;
        T gamma_;
        std::mt19937 rng_;

        nn::Linear<T> net1_, net2_;
        nn::ReLU<T> relu_;
        nn::Softmax<T> softmax_;
        nn::Adam<T> optimizer_;

        std::vector<T> lastProbs_;
        std::vector<std::vector<T>> states_;
        std::vector<size_t> actions_;
        std::vector<T> rewards_;
    };

} // namespace calx::ml

#endif // CALX_ML_RL_POLICYGRADIENT_HPP
