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

// DQN.hpp — Deep Q-Network
//
// NN で Q(s,a) を近似。Experience Replay + Target Network。

#ifndef CALX_ML_RL_DQN_HPP
#define CALX_ML_RL_DQN_HPP

#include <ml/nn/Layer.hpp>
#include <ml/nn/Activation.hpp>
#include <ml/nn/Loss.hpp>
#include <ml/nn/Optimizer.hpp>
#include <algorithm>
#include <deque>
#include <random>
#include <tuple>

namespace calx::ml {

    /// Experience Replay バッファ
    template<typename T>
    struct Experience {
        std::vector<T> state;
        size_t action;
        T reward;
        std::vector<T> nextState;
        bool done;
    };

    /// DQN (Deep Q-Network)
    template<typename T>
    class DQN {
    public:
        /// @param stateDim    状態次元
        /// @param numActions  行動数
        /// @param hiddenDim   隠れ層次元
        /// @param lr          学習率
        /// @param gamma       割引率
        /// @param epsilon     ε-greedy の探索率
        /// @param bufferSize  Experience Replay バッファサイズ
        /// @param batchSize   ミニバッチサイズ
        /// @param seed        乱数シード
        DQN(size_t stateDim, size_t numActions, size_t hiddenDim = 64,
            T lr = T{0.001}, T gamma = T{0.99}, T epsilon = T{1.0},
            size_t bufferSize = 10000, size_t batchSize = 32, unsigned seed = 42)
            : stateDim_(stateDim), nActions_(numActions)
            , gamma_(gamma), epsilon_(epsilon)
            , bufferSize_(bufferSize), batchSize_(batchSize), rng_(seed)
            // Q ネットワーク
            , qNet1_(stateDim, hiddenDim, true, seed)
            , qNet2_(hiddenDim, numActions, true, seed + 1)
            , relu_()
            // ターゲットネットワーク (同じ構造)
            , tNet1_(stateDim, hiddenDim, true, seed)
            , tNet2_(hiddenDim, numActions, true, seed + 1)
            , tRelu_()
            , optimizer_(lr)
        {
            updateTargetNetwork();
        }

        /// ε-greedy で行動選択
        size_t selectAction(const std::vector<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(const std::vector<T>& state) {
            Matrix<T> input(1, stateDim_);
            for (size_t d = 0; d < stateDim_; ++d) input(0, d) = state[d];
            auto qvals = forwardQ(input);
            size_t bestA = 0;
            T bestQ = qvals(0, 0);
            for (size_t a = 1; a < nActions_; ++a)
                if (qvals(0, a) > bestQ) { bestQ = qvals(0, a); bestA = a; }
            return bestA;
        }

        /// 経験をバッファに追加
        void store(const std::vector<T>& state, size_t action, T reward,
                   const std::vector<T>& nextState, bool done) {
            if (buffer_.size() >= bufferSize_) buffer_.pop_front();
            buffer_.push_back({ state, action, reward, nextState, done });
        }

        /// ミニバッチ学習 (1ステップ)
        /// @return 損失値 (バッファ不足時は -1)
        T trainStep() {
            if (buffer_.size() < batchSize_) return T{-1};

            // ランダムサンプリング
            std::vector<size_t> indices(buffer_.size());
            std::iota(indices.begin(), indices.end(), 0);
            std::shuffle(indices.begin(), indices.end(), rng_);

            Matrix<T> states(batchSize_, stateDim_);
            Matrix<T> nextStates(batchSize_, stateDim_);
            std::vector<size_t> actions(batchSize_);
            std::vector<T> rewards(batchSize_);
            std::vector<bool> dones(batchSize_);

            for (size_t i = 0; i < batchSize_; ++i) {
                auto& exp = buffer_[indices[i]];
                for (size_t d = 0; d < stateDim_; ++d) {
                    states(i, d) = exp.state[d];
                    nextStates(i, d) = exp.nextState[d];
                }
                actions[i] = exp.action;
                rewards[i] = exp.reward;
                dones[i] = exp.done;
            }

            // Q 値計算
            auto qvals = forwardQ(states);

            // ターゲット Q 値
            auto nextQ = forwardTarget(nextStates);

            // ターゲット構築
            Matrix<T> target = qvals;  // コピー
            for (size_t i = 0; i < batchSize_; ++i) {
                T maxNextQ = T{0};
                if (!dones[i]) {
                    maxNextQ = nextQ(i, 0);
                    for (size_t a = 1; a < nActions_; ++a)
                        if (nextQ(i, a) > maxNextQ) maxNextQ = nextQ(i, a);
                }
                target(i, actions[i]) = rewards[i] + gamma_ * maxNextQ;
            }

            // MSE 損失
            Matrix<T> grad;
            T loss = nn::mseLoss(qvals, target, grad);

            // 逆伝播
            auto params = qParameters();
            nn::OptimizerBase<T>::zeroGrad(params);
            auto d2 = qNet2_.backward(grad);
            qNet1_.backward(relu_.backward(d2));
            optimizer_.step(params);

            return loss;
        }

        /// ターゲットネットワークを更新 (ハードコピー)
        void updateTargetNetwork() {
            copyWeights(qNet1_, tNet1_);
            copyWeights(qNet2_, tNet2_);
        }

        void setEpsilon(T eps) { epsilon_ = eps; }
        size_t bufferSize() const { return buffer_.size(); }

    private:
        Matrix<T> forwardQ(const Matrix<T>& input) {
            return qNet2_.forward(relu_.forward(qNet1_.forward(input)));
        }

        Matrix<T> forwardTarget(const Matrix<T>& input) {
            return tNet2_.forward(tRelu_.forward(tNet1_.forward(input)));
        }

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

        static void copyWeights(nn::Linear<T>& src, nn::Linear<T>& dst) {
            auto sp = src.parameters();
            auto dp = dst.parameters();
            for (size_t i = 0; i < sp.size(); ++i) {
                auto& sw = *sp[i].param;
                auto& dw = *dp[i].param;
                for (size_t r = 0; r < sw.rows(); ++r)
                    for (size_t c = 0; c < sw.cols(); ++c)
                        dw(r, c) = sw(r, c);
            }
        }

        size_t stateDim_, nActions_;
        T gamma_, epsilon_;
        size_t bufferSize_, batchSize_;
        std::mt19937 rng_;

        nn::Linear<T> qNet1_, qNet2_;
        nn::ReLU<T> relu_;
        nn::Linear<T> tNet1_, tNet2_;
        nn::ReLU<T> tRelu_;
        nn::Adam<T> optimizer_;

        std::deque<Experience<T>> buffer_;
    };

} // namespace calx::ml

#endif // CALX_ML_RL_DQN_HPP
