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

// RNN.hpp — リカレントニューラルネットワーク
//
// RNN (Elman), LSTM, GRU。
// 入力: Matrix (batch × seqLen*inputSize), 出力: Matrix (batch × seqLen*hiddenSize)
// 最終隠れ状態のみ必要な場合は lastHidden() を使用。

#ifndef CALX_ML_NN_RNN_HPP
#define CALX_ML_NN_RNN_HPP

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

namespace calx::ml::nn {

    // =================================================================
    // RNN (Elman RNN)
    // =================================================================

    /// h_t = tanh(W_ih * x_t + W_hh * h_{t-1} + b)
    template<typename T>
    class RNNCell : public Layer<T> {
    public:
        RNNCell(size_t inputSize, size_t hiddenSize, size_t seqLen,
                unsigned seed = 42)
            : inSize_(inputSize), hSize_(hiddenSize), seqLen_(seqLen)
        {
            T scale = std::sqrt(T{1} / static_cast<T>(hiddenSize));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            Wih_ = Matrix<T>(hSize_, inSize_);
            Whh_ = Matrix<T>(hSize_, hSize_);
            bh_ = Matrix<T>(1, hSize_, T{0});
            for (size_t i = 0; i < hSize_; ++i) {
                for (size_t j = 0; j < inSize_; ++j) Wih_(i, j) = dist(rng);
                for (size_t j = 0; j < hSize_; ++j) Whh_(i, j) = dist(rng);
            }
            gWih_ = Matrix<T>(hSize_, inSize_, T{0});
            gWhh_ = Matrix<T>(hSize_, hSize_, T{0});
            gbh_ = Matrix<T>(1, hSize_, T{0});
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            // 隠れ状態を保存 (BPTT 用)
            hStates_.resize(seqLen_ + 1);
            for (auto& h : hStates_) h.assign(batch * hSize_, T{0});
            xCache_.resize(seqLen_);

            Matrix<T> output(batch, seqLen_ * hSize_);

            for (size_t t = 0; t < seqLen_; ++t) {
                xCache_[t].resize(batch * inSize_);
                for (size_t b = 0; b < batch; ++b)
                    for (size_t i = 0; i < inSize_; ++i)
                        xCache_[t][b * inSize_ + i] = input(b, t * inSize_ + i);

                for (size_t b = 0; b < batch; ++b) {
                    for (size_t h = 0; h < hSize_; ++h) {
                        T sum = bh_(0, h);
                        for (size_t i = 0; i < inSize_; ++i)
                            sum += Wih_(h, i) * xCache_[t][b * inSize_ + i];
                        for (size_t j = 0; j < hSize_; ++j)
                            sum += Whh_(h, j) * hStates_[t][b * hSize_ + j];
                        hStates_[t + 1][b * hSize_ + h] = std::tanh(sum);
                        output(b, t * hSize_ + h) = hStates_[t + 1][b * hSize_ + h];
                    }
                }
            }
            batch_ = batch;
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = batch_;
            Matrix<T> gradInput(batch, seqLen_ * inSize_, T{0});
            std::vector<T> dh(batch * hSize_, T{0});

            for (int t = static_cast<int>(seqLen_) - 1; t >= 0; --t) {
                for (size_t b = 0; b < batch; ++b) {
                    for (size_t h = 0; h < hSize_; ++h) {
                        T go = gradOutput(b, t * hSize_ + h) + dh[b * hSize_ + h];
                        T ht = hStates_[t + 1][b * hSize_ + h];
                        T dtanh = go * (T{1} - ht * ht);

                        // gradients
                        for (size_t i = 0; i < inSize_; ++i) {
                            gWih_(h, i) += dtanh * xCache_[t][b * inSize_ + i];
                            gradInput(b, t * inSize_ + i) += dtanh * Wih_(h, i);
                        }
                        for (size_t j = 0; j < hSize_; ++j)
                            gWhh_(h, j) += dtanh * hStates_[t][b * hSize_ + j];
                        gbh_(0, h) += dtanh;

                        // dh_{t-1}
                        dh[b * hSize_ + h] = T{0};
                    }
                    // accumulate dh for previous step
                    for (size_t h = 0; h < hSize_; ++h) {
                        T go = gradOutput(b, t * hSize_ + h) + (t < static_cast<int>(seqLen_) - 1 ? T{0} : T{0});
                        T ht = hStates_[t + 1][b * hSize_ + h];
                        T dtanh = (gradOutput(b, t * hSize_ + h) +
                                   (t == static_cast<int>(seqLen_) - 1 ? T{0} : dh[b * hSize_ + h]))
                                  * (T{1} - ht * ht);
                        for (size_t j = 0; j < hSize_; ++j)
                            dh[b * hSize_ + j] += dtanh * Whh_(h, j);
                    }
                    // Reset dh after accumulation (it now holds dh for step t-1)
                }
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&Wih_, &gWih_}, {&Whh_, &gWhh_}, {&bh_, &gbh_} };
        }

    private:
        size_t inSize_, hSize_, seqLen_, batch_ = 0;
        Matrix<T> Wih_, Whh_, bh_;
        Matrix<T> gWih_, gWhh_, gbh_;
        std::vector<std::vector<T>> hStates_, xCache_;
    };

    // =================================================================
    // LSTM
    // =================================================================

    template<typename T>
    class LSTMCell : public Layer<T> {
    public:
        LSTMCell(size_t inputSize, size_t hiddenSize, size_t seqLen,
                 unsigned seed = 42)
            : inSize_(inputSize), hSize_(hiddenSize), seqLen_(seqLen)
        {
            T scale = std::sqrt(T{1} / static_cast<T>(hiddenSize));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            size_t gateSize = 4 * hSize_;  // i, f, g, o gates
            Wi_ = Matrix<T>(gateSize, inSize_);
            Wh_ = Matrix<T>(gateSize, hSize_);
            bias_ = Matrix<T>(1, gateSize, T{0});
            for (size_t i = 0; i < gateSize; ++i) {
                for (size_t j = 0; j < inSize_; ++j) Wi_(i, j) = dist(rng);
                for (size_t j = 0; j < hSize_; ++j) Wh_(i, j) = dist(rng);
            }
            // forget gate bias を 1 に初期化
            for (size_t h = 0; h < hSize_; ++h)
                bias_(0, hSize_ + h) = T{1};

            gWi_ = Matrix<T>(gateSize, inSize_, T{0});
            gWh_ = Matrix<T>(gateSize, hSize_, T{0});
            gBias_ = Matrix<T>(1, gateSize, T{0});
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            batch_ = batch;
            size_t gs = 4 * hSize_;

            // cache all states
            gates_.resize(seqLen_);
            hCache_.resize(seqLen_ + 1);
            cCache_.resize(seqLen_ + 1);
            xCache_.resize(seqLen_);
            for (auto& v : hCache_) v.assign(batch * hSize_, T{0});
            for (auto& v : cCache_) v.assign(batch * hSize_, T{0});

            Matrix<T> output(batch, seqLen_ * hSize_);

            for (size_t t = 0; t < seqLen_; ++t) {
                xCache_[t].resize(batch * inSize_);
                for (size_t b = 0; b < batch; ++b)
                    for (size_t i = 0; i < inSize_; ++i)
                        xCache_[t][b * inSize_ + i] = input(b, t * inSize_ + i);

                gates_[t].resize(batch * gs);

                for (size_t b = 0; b < batch; ++b) {
                    // Compute all 4 gates
                    for (size_t g = 0; g < gs; ++g) {
                        T sum = bias_(0, g);
                        for (size_t i = 0; i < inSize_; ++i)
                            sum += Wi_(g, i) * xCache_[t][b * inSize_ + i];
                        for (size_t h = 0; h < hSize_; ++h)
                            sum += Wh_(g, h) * hCache_[t][b * hSize_ + h];
                        gates_[t][b * gs + g] = sum;
                    }

                    // Apply activations: i=sigmoid, f=sigmoid, g=tanh, o=sigmoid
                    for (size_t h = 0; h < hSize_; ++h) {
                        T ig = sig(gates_[t][b * gs + h]);                    // input gate
                        T fg = sig(gates_[t][b * gs + hSize_ + h]);           // forget gate
                        T gg = std::tanh(gates_[t][b * gs + 2 * hSize_ + h]); // cell gate
                        T og = sig(gates_[t][b * gs + 3 * hSize_ + h]);       // output gate

                        // Store activated gates
                        gates_[t][b * gs + h] = ig;
                        gates_[t][b * gs + hSize_ + h] = fg;
                        gates_[t][b * gs + 2 * hSize_ + h] = gg;
                        gates_[t][b * gs + 3 * hSize_ + h] = og;

                        T ct = fg * cCache_[t][b * hSize_ + h] + ig * gg;
                        T ht = og * std::tanh(ct);

                        cCache_[t + 1][b * hSize_ + h] = ct;
                        hCache_[t + 1][b * hSize_ + h] = ht;
                        output(b, t * hSize_ + h) = ht;
                    }
                }
            }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = batch_;
            size_t gs = 4 * hSize_;
            Matrix<T> gradInput(batch, seqLen_ * inSize_, T{0});
            std::vector<T> dh(batch * hSize_, T{0});
            std::vector<T> dc(batch * hSize_, T{0});

            for (int t = static_cast<int>(seqLen_) - 1; t >= 0; --t) {
                for (size_t b = 0; b < batch; ++b) {
                    for (size_t h = 0; h < hSize_; ++h) {
                        T go = gradOutput(b, t * hSize_ + h) + dh[b * hSize_ + h];
                        T ig = gates_[t][b * gs + h];
                        T fg = gates_[t][b * gs + hSize_ + h];
                        T gg = gates_[t][b * gs + 2 * hSize_ + h];
                        T og = gates_[t][b * gs + 3 * hSize_ + h];
                        T ct = cCache_[t + 1][b * hSize_ + h];
                        T tanhCt = std::tanh(ct);

                        T dOg = go * tanhCt * og * (T{1} - og);
                        T dCt = go * og * (T{1} - tanhCt * tanhCt) + dc[b * hSize_ + h];
                        T dIg = dCt * gg * ig * (T{1} - ig);
                        T dFg = dCt * cCache_[t][b * hSize_ + h] * fg * (T{1} - fg);
                        T dGg = dCt * ig * (T{1} - gg * gg);

                        dc[b * hSize_ + h] = dCt * fg;

                        // Pack gate gradients
                        T dGates[4] = { dIg, dFg, dGg, dOg };
                        for (int g = 0; g < 4; ++g) {
                            size_t offset = g * hSize_ + h;
                            gBias_(0, offset) += dGates[g];
                            for (size_t i = 0; i < inSize_; ++i) {
                                gWi_(offset, i) += dGates[g] * xCache_[t][b * inSize_ + i];
                                gradInput(b, t * inSize_ + i) += dGates[g] * Wi_(offset, i);
                            }
                            for (size_t j = 0; j < hSize_; ++j) {
                                gWh_(offset, j) += dGates[g] * hCache_[t][b * hSize_ + j];
                                dh[b * hSize_ + j] += dGates[g] * Wh_(offset, j);
                            }
                        }
                    }
                }
                // dh now holds gradient for h_{t-1}
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&Wi_, &gWi_}, {&Wh_, &gWh_}, {&bias_, &gBias_} };
        }

    private:
        static T sig(T x) { return T{1} / (T{1} + std::exp(-x)); }

        size_t inSize_, hSize_, seqLen_, batch_ = 0;
        Matrix<T> Wi_, Wh_, bias_;
        Matrix<T> gWi_, gWh_, gBias_;
        std::vector<std::vector<T>> gates_, hCache_, cCache_, xCache_;
    };

    // =================================================================
    // GRU
    // =================================================================

    template<typename T>
    class GRUCell : public Layer<T> {
    public:
        GRUCell(size_t inputSize, size_t hiddenSize, size_t seqLen,
                unsigned seed = 42)
            : inSize_(inputSize), hSize_(hiddenSize), seqLen_(seqLen)
        {
            T scale = std::sqrt(T{1} / static_cast<T>(hiddenSize));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            size_t gs = 3 * hSize_;  // r, z, n gates
            Wi_ = Matrix<T>(gs, inSize_);
            Wh_ = Matrix<T>(gs, hSize_);
            bias_ = Matrix<T>(1, gs, T{0});
            for (size_t i = 0; i < gs; ++i) {
                for (size_t j = 0; j < inSize_; ++j) Wi_(i, j) = dist(rng);
                for (size_t j = 0; j < hSize_; ++j) Wh_(i, j) = dist(rng);
            }
            gWi_ = Matrix<T>(gs, inSize_, T{0});
            gWh_ = Matrix<T>(gs, hSize_, T{0});
            gBias_ = Matrix<T>(1, gs, T{0});
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            batch_ = batch;

            gates_.resize(seqLen_);
            hCache_.resize(seqLen_ + 1);
            xCache_.resize(seqLen_);
            for (auto& v : hCache_) v.assign(batch * hSize_, T{0});

            Matrix<T> output(batch, seqLen_ * hSize_);

            for (size_t t = 0; t < seqLen_; ++t) {
                xCache_[t].resize(batch * inSize_);
                for (size_t b = 0; b < batch; ++b)
                    for (size_t i = 0; i < inSize_; ++i)
                        xCache_[t][b * inSize_ + i] = input(b, t * inSize_ + i);

                gates_[t].resize(batch * 3 * hSize_);

                for (size_t b = 0; b < batch; ++b) {
                    // r, z gates
                    for (size_t h = 0; h < hSize_; ++h) {
                        T rSum = bias_(0, h), zSum = bias_(0, hSize_ + h);
                        for (size_t i = 0; i < inSize_; ++i) {
                            rSum += Wi_(h, i) * xCache_[t][b * inSize_ + i];
                            zSum += Wi_(hSize_ + h, i) * xCache_[t][b * inSize_ + i];
                        }
                        for (size_t j = 0; j < hSize_; ++j) {
                            rSum += Wh_(h, j) * hCache_[t][b * hSize_ + j];
                            zSum += Wh_(hSize_ + j, j) * hCache_[t][b * hSize_ + j];
                        }
                        T r = sig(rSum), z = sig(zSum);
                        gates_[t][b * 3 * hSize_ + h] = r;
                        gates_[t][b * 3 * hSize_ + hSize_ + h] = z;

                        // n gate (new)
                        T nSum = bias_(0, 2 * hSize_ + h);
                        for (size_t i = 0; i < inSize_; ++i)
                            nSum += Wi_(2 * hSize_ + h, i) * xCache_[t][b * inSize_ + i];
                        for (size_t j = 0; j < hSize_; ++j)
                            nSum += Wh_(2 * hSize_ + h, j) * r * hCache_[t][b * hSize_ + j];
                        T n = std::tanh(nSum);
                        gates_[t][b * 3 * hSize_ + 2 * hSize_ + h] = n;

                        T ht = (T{1} - z) * n + z * hCache_[t][b * hSize_ + h];
                        hCache_[t + 1][b * hSize_ + h] = ht;
                        output(b, t * hSize_ + h) = ht;
                    }
                }
            }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = batch_;
            Matrix<T> gradInput(batch, seqLen_ * inSize_, T{0});
            std::vector<T> dh(batch * hSize_, T{0});

            for (int t = static_cast<int>(seqLen_) - 1; t >= 0; --t) {
                for (size_t b = 0; b < batch; ++b) {
                    for (size_t h = 0; h < hSize_; ++h) {
                        T go = gradOutput(b, t * hSize_ + h) + dh[b * hSize_ + h];
                        T r = gates_[t][b * 3 * hSize_ + h];
                        T z = gates_[t][b * 3 * hSize_ + hSize_ + h];
                        T n = gates_[t][b * 3 * hSize_ + 2 * hSize_ + h];
                        T hPrev = hCache_[t][b * hSize_ + h];

                        T dz = go * (hPrev - n) * z * (T{1} - z);
                        T dn = go * (T{1} - z) * (T{1} - n * n);

                        // gradients for z gate
                        gBias_(0, hSize_ + h) += dz;
                        for (size_t i = 0; i < inSize_; ++i) {
                            gWi_(hSize_ + h, i) += dz * xCache_[t][b * inSize_ + i];
                            gradInput(b, t * inSize_ + i) += dz * Wi_(hSize_ + h, i);
                        }

                        // gradients for n gate
                        gBias_(0, 2 * hSize_ + h) += dn;
                        for (size_t i = 0; i < inSize_; ++i) {
                            gWi_(2 * hSize_ + h, i) += dn * xCache_[t][b * inSize_ + i];
                            gradInput(b, t * inSize_ + i) += dn * Wi_(2 * hSize_ + h, i);
                        }
                        for (size_t j = 0; j < hSize_; ++j)
                            gWh_(2 * hSize_ + h, j) += dn * r * hCache_[t][b * hSize_ + j];

                        // dr from n gate
                        T drFromN = T{0};
                        for (size_t j = 0; j < hSize_; ++j)
                            drFromN += dn * Wh_(2 * hSize_ + h, j) * hCache_[t][b * hSize_ + j];
                        T dr = drFromN * r * (T{1} - r);

                        // gradients for r gate
                        gBias_(0, h) += dr;
                        for (size_t i = 0; i < inSize_; ++i) {
                            gWi_(h, i) += dr * xCache_[t][b * inSize_ + i];
                            gradInput(b, t * inSize_ + i) += dr * Wi_(h, i);
                        }

                        // dh_{t-1}
                        dh[b * hSize_ + h] = go * z;
                        for (size_t j = 0; j < hSize_; ++j) {
                            dh[b * hSize_ + j] += dr * Wh_(h, j);
                            dh[b * hSize_ + j] += dz * Wh_(hSize_ + h, j);
                        }
                    }
                }
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&Wi_, &gWi_}, {&Wh_, &gWh_}, {&bias_, &gBias_} };
        }

    private:
        static T sig(T x) { return T{1} / (T{1} + std::exp(-x)); }

        size_t inSize_, hSize_, seqLen_, batch_ = 0;
        Matrix<T> Wi_, Wh_, bias_;
        Matrix<T> gWi_, gWh_, gBias_;
        std::vector<std::vector<T>> gates_, hCache_, xCache_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_RNN_HPP
