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

// Transformer.hpp — Transformer ブロック
//
// TransformerEncoderLayer: MultiHeadAttention + FFN + LayerNorm + Residual
// TransformerDecoderLayer: Self-Attn + Cross-Attn + FFN + LayerNorm + Residual
// TransformerEncoder / TransformerDecoder: N 層スタック

#ifndef CALX_ML_NN_TRANSFORMER_HPP
#define CALX_ML_NN_TRANSFORMER_HPP

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

namespace calx::ml::nn {

    // =================================================================
    // LayerNorm (層正規化)
    // =================================================================

    /// 入力: (batch × seqLen*dModel), 各位置の dModel 次元ベクトルを正規化
    template<typename T>
    class LayerNorm : public Layer<T> {
    public:
        LayerNorm(size_t dModel, size_t seqLen, T eps = T{1e-5})
            : dModel_(dModel), seqLen_(seqLen), eps_(eps)
        {
            gamma_ = Matrix<T>(1, dModel_, T{1});
            beta_ = Matrix<T>(1, dModel_, T{0});
            gGamma_ = Matrix<T>(1, dModel_, T{0});
            gBeta_ = Matrix<T>(1, dModel_, T{0});
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            xNorm_ = Matrix<T>(batch, seqLen_ * dModel_);
            mean_.resize(batch * seqLen_);
            invStd_.resize(batch * seqLen_);

            Matrix<T> output(batch, seqLen_ * dModel_);
            for (size_t b = 0; b < batch; ++b) {
                for (size_t s = 0; s < seqLen_; ++s) {
                    size_t base = s * dModel_;
                    size_t idx = b * seqLen_ + s;

                    T m = T{0};
                    for (size_t d = 0; d < dModel_; ++d)
                        m += input(b, base + d);
                    m /= static_cast<T>(dModel_);
                    mean_[idx] = m;

                    T v = T{0};
                    for (size_t d = 0; d < dModel_; ++d) {
                        T diff = input(b, base + d) - m;
                        v += diff * diff;
                    }
                    v /= static_cast<T>(dModel_);
                    invStd_[idx] = T{1} / std::sqrt(v + eps_);

                    for (size_t d = 0; d < dModel_; ++d) {
                        T xn = (input(b, base + d) - m) * invStd_[idx];
                        xNorm_(b, base + d) = xn;
                        output(b, base + d) = gamma_(0, d) * xn + beta_(0, d);
                    }
                }
            }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = gradOutput.rows();
            Matrix<T> gradInput(batch, seqLen_ * dModel_);
            T invD = T{1} / static_cast<T>(dModel_);

            for (size_t b = 0; b < batch; ++b) {
                for (size_t s = 0; s < seqLen_; ++s) {
                    size_t base = s * dModel_;
                    size_t idx = b * seqLen_ + s;

                    // γ, β の勾配を累積
                    for (size_t d = 0; d < dModel_; ++d) {
                        gGamma_(0, d) += gradOutput(b, base + d) * xNorm_(b, base + d);
                        gBeta_(0, d) += gradOutput(b, base + d);
                    }

                    // 入力勾配
                    T dxnSum = T{0}, dxnXnSum = T{0};
                    for (size_t d = 0; d < dModel_; ++d) {
                        T dxn = gradOutput(b, base + d) * gamma_(0, d);
                        dxnSum += dxn;
                        dxnXnSum += dxn * xNorm_(b, base + d);
                    }
                    for (size_t d = 0; d < dModel_; ++d) {
                        T dxn = gradOutput(b, base + d) * gamma_(0, d);
                        gradInput(b, base + d) = invStd_[idx] * invD *
                            (static_cast<T>(dModel_) * dxn - dxnSum
                             - xNorm_(b, base + d) * dxnXnSum);
                    }
                }
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&gamma_, &gGamma_}, {&beta_, &gBeta_} };
        }

    private:
        size_t dModel_, seqLen_;
        T eps_;
        Matrix<T> gamma_, beta_, gGamma_, gBeta_;
        Matrix<T> xNorm_;
        std::vector<T> mean_, invStd_;
    };

    // =================================================================
    // FeedForward (Position-wise FFN)
    // =================================================================

    /// FFN(x) = ReLU(x*W1 + b1) * W2 + b2
    template<typename T>
    class FeedForward : public Layer<T> {
    public:
        FeedForward(size_t dModel, size_t dFF, size_t seqLen, unsigned seed = 42)
            : dModel_(dModel), dFF_(dFF), seqLen_(seqLen)
        {
            T scale1 = std::sqrt(T{2} / static_cast<T>(dModel));
            T scale2 = std::sqrt(T{1} / static_cast<T>(dFF));
            std::mt19937 rng(seed);
            std::normal_distribution<T> d1(T{0}, scale1), d2(T{0}, scale2);

            W1_ = Matrix<T>(dFF, dModel);
            b1_ = Matrix<T>(1, dFF, T{0});
            W2_ = Matrix<T>(dModel, dFF);
            b2_ = Matrix<T>(1, dModel, T{0});
            for (size_t i = 0; i < dFF; ++i)
                for (size_t j = 0; j < dModel; ++j) W1_(i, j) = d1(rng);
            for (size_t i = 0; i < dModel; ++i)
                for (size_t j = 0; j < dFF; ++j) W2_(i, j) = d2(rng);

            gW1_ = Matrix<T>(dFF, dModel, T{0});
            gb1_ = Matrix<T>(1, dFF, T{0});
            gW2_ = Matrix<T>(dModel, dFF, T{0});
            gb2_ = Matrix<T>(1, dModel, T{0});
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            size_t total = seqLen_ * dModel_;
            hidden_ = Matrix<T>(batch * seqLen_, dFF_);
            Matrix<T> output(batch, total);

            for (size_t b = 0; b < batch; ++b) {
                for (size_t s = 0; s < seqLen_; ++s) {
                    size_t row = b * seqLen_ + s;
                    // h = ReLU(x * W1^T + b1)
                    for (size_t f = 0; f < dFF_; ++f) {
                        T sum = b1_(0, f);
                        for (size_t d = 0; d < dModel_; ++d)
                            sum += input(b, s * dModel_ + d) * W1_(f, d);
                        hidden_(row, f) = sum > T{0} ? sum : T{0};
                    }
                    // out = h * W2^T + b2
                    for (size_t d = 0; d < dModel_; ++d) {
                        T sum = b2_(0, d);
                        for (size_t f = 0; f < dFF_; ++f)
                            sum += hidden_(row, f) * W2_(d, f);
                        output(b, s * dModel_ + d) = sum;
                    }
                }
            }
            input_ = input;
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = gradOutput.rows();
            Matrix<T> gradInput(batch, seqLen_ * dModel_, T{0});

            for (size_t b = 0; b < batch; ++b) {
                for (size_t s = 0; s < seqLen_; ++s) {
                    size_t row = b * seqLen_ + s;
                    // dh = gradOut * W2
                    for (size_t f = 0; f < dFF_; ++f) {
                        T dh = T{0};
                        for (size_t d = 0; d < dModel_; ++d) {
                            T go = gradOutput(b, s * dModel_ + d);
                            dh += go * W2_(d, f);
                            gW2_(d, f) += go * hidden_(row, f);
                        }
                        // ReLU backward
                        dh = hidden_(row, f) > T{0} ? dh : T{0};
                        // dx = dh * W1
                        for (size_t d = 0; d < dModel_; ++d) {
                            gradInput(b, s * dModel_ + d) += dh * W1_(f, d);
                            gW1_(f, d) += dh * input_(b, s * dModel_ + d);
                        }
                        gb1_(0, f) += dh;
                    }
                    for (size_t d = 0; d < dModel_; ++d)
                        gb2_(0, d) += gradOutput(b, s * dModel_ + d);
                }
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&W1_, &gW1_}, {&b1_, &gb1_}, {&W2_, &gW2_}, {&b2_, &gb2_} };
        }

    private:
        size_t dModel_, dFF_, seqLen_;
        Matrix<T> W1_, b1_, W2_, b2_;
        Matrix<T> gW1_, gb1_, gW2_, gb2_;
        Matrix<T> hidden_, input_;
    };

    // =================================================================
    // Transformer Encoder Layer
    // =================================================================

    /// Self-Attention + FFN + LayerNorm + Residual
    template<typename T>
    class TransformerEncoderLayer : public Layer<T> {
    public:
        TransformerEncoderLayer(size_t dModel, size_t numHeads, size_t dFF,
                                 size_t seqLen, unsigned seed = 42)
            : attn_(dModel, numHeads, seqLen, seed)
            , norm1_(dModel, seqLen)
            , ffn_(dModel, dFF, seqLen, seed + 1)
            , norm2_(dModel, seqLen)
        {}

        Matrix<T> forward(const Matrix<T>& input) override {
            input1_ = input;
            // Sub-layer 1: MultiHeadAttention + Residual + LayerNorm
            auto attnOut = attn_.forward(input);
            auto res1 = addResidual(input, attnOut);
            auto norm1Out = norm1_.forward(res1);

            input2_ = norm1Out;
            // Sub-layer 2: FFN + Residual + LayerNorm
            auto ffnOut = ffn_.forward(norm1Out);
            auto res2 = addResidual(norm1Out, ffnOut);
            return norm2_.forward(res2);
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            // norm2 backward
            auto dRes2 = norm2_.backward(gradOutput);
            // residual: grad flows to both ffn and skip
            auto dFFN = ffn_.backward(dRes2);
            auto dNorm1 = addResidual(dRes2, dFFN);

            // norm1 backward
            auto dRes1 = norm1_.backward(dNorm1);
            // residual: grad flows to both attn and skip
            auto dAttn = attn_.backward(dRes1);
            return addResidual(dRes1, dAttn);
        }

        std::vector<ParamGrad<T>> parameters() override {
            auto p1 = attn_.parameters();
            auto p2 = norm1_.parameters();
            auto p3 = ffn_.parameters();
            auto p4 = norm2_.parameters();
            p1.insert(p1.end(), p2.begin(), p2.end());
            p1.insert(p1.end(), p3.begin(), p3.end());
            p1.insert(p1.end(), p4.begin(), p4.end());
            return p1;
        }

    private:
        static Matrix<T> addResidual(const Matrix<T>& a, const Matrix<T>& b) {
            Matrix<T> result(a.rows(), a.cols());
            for (size_t i = 0; i < a.rows(); ++i)
                for (size_t j = 0; j < a.cols(); ++j)
                    result(i, j) = a(i, j) + b(i, j);
            return result;
        }

        MultiHeadAttention<T> attn_;
        LayerNorm<T> norm1_;
        FeedForward<T> ffn_;
        LayerNorm<T> norm2_;
        Matrix<T> input1_, input2_;
    };

    // =================================================================
    // Transformer Decoder Layer
    // =================================================================

    /// Self-Attn + Cross-Attn + FFN + LayerNorm + Residual
    template<typename T>
    class TransformerDecoderLayer : public Layer<T> {
    public:
        TransformerDecoderLayer(size_t dModel, size_t numHeads, size_t dFF,
                                 size_t seqLen, unsigned seed = 42)
            : selfAttn_(dModel, numHeads, seqLen, seed)
            , crossAttn_(dModel, numHeads, seqLen, seed + 1)
            , norm1_(dModel, seqLen)
            , norm2_(dModel, seqLen)
            , norm3_(dModel, seqLen)
            , ffn_(dModel, dFF, seqLen, seed + 2)
        {}

        /// エンコーダ出力をセット (cross-attention 用)
        void setEncoderOutput(const Matrix<T>& encOut) { encOut_ = encOut; }

        Matrix<T> forward(const Matrix<T>& input) override {
            // Sub-layer 1: Self-Attention + Residual + LayerNorm
            auto selfAttnOut = selfAttn_.forward(input);
            auto res1 = addResidual(input, selfAttnOut);
            auto norm1Out = norm1_.forward(res1);

            // Sub-layer 2: Cross-Attention (query=decoder, key/value=encoder)
            // 簡略: encOut を入力として cross-attention
            auto crossOut = crossAttn_.forward(norm1Out);  // 実際は enc を K,V に使うべき
            auto res2 = addResidual(norm1Out, crossOut);
            auto norm2Out = norm2_.forward(res2);

            // Sub-layer 3: FFN + Residual + LayerNorm
            auto ffnOut = ffn_.forward(norm2Out);
            auto res3 = addResidual(norm2Out, ffnOut);
            return norm3_.forward(res3);
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            auto dRes3 = norm3_.backward(gradOutput);
            auto dFFN = ffn_.backward(dRes3);
            auto dNorm2 = addResidual(dRes3, dFFN);
            auto dRes2 = norm2_.backward(dNorm2);
            auto dCross = crossAttn_.backward(dRes2);
            auto dNorm1 = addResidual(dRes2, dCross);
            auto dRes1 = norm1_.backward(dNorm1);
            auto dSelf = selfAttn_.backward(dRes1);
            return addResidual(dRes1, dSelf);
        }

        std::vector<ParamGrad<T>> parameters() override {
            auto p = selfAttn_.parameters();
            auto p2 = crossAttn_.parameters();
            auto p3 = norm1_.parameters();
            auto p4 = norm2_.parameters();
            auto p5 = norm3_.parameters();
            auto p6 = ffn_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            p.insert(p.end(), p3.begin(), p3.end());
            p.insert(p.end(), p4.begin(), p4.end());
            p.insert(p.end(), p5.begin(), p5.end());
            p.insert(p.end(), p6.begin(), p6.end());
            return p;
        }

    private:
        static Matrix<T> addResidual(const Matrix<T>& a, const Matrix<T>& b) {
            Matrix<T> result(a.rows(), a.cols());
            for (size_t i = 0; i < a.rows(); ++i)
                for (size_t j = 0; j < a.cols(); ++j)
                    result(i, j) = a(i, j) + b(i, j);
            return result;
        }

        MultiHeadAttention<T> selfAttn_, crossAttn_;
        LayerNorm<T> norm1_, norm2_, norm3_;
        FeedForward<T> ffn_;
        Matrix<T> encOut_;
    };

    // =================================================================
    // Transformer Encoder (N 層スタック)
    // =================================================================

    template<typename T>
    class TransformerEncoder {
    public:
        TransformerEncoder(size_t numLayers, size_t dModel, size_t numHeads,
                            size_t dFF, size_t seqLen, unsigned seed = 42)
        {
            for (size_t i = 0; i < numLayers; ++i)
                layers_.push_back(std::make_unique<TransformerEncoderLayer<T>>(
                    dModel, numHeads, dFF, seqLen, seed + static_cast<unsigned>(i * 10)));
        }

        Matrix<T> forward(const Matrix<T>& input) {
            Matrix<T> x = input;
            for (auto& layer : layers_) x = layer->forward(x);
            return x;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) {
            Matrix<T> grad = gradOutput;
            for (int i = static_cast<int>(layers_.size()) - 1; i >= 0; --i)
                grad = layers_[i]->backward(grad);
            return grad;
        }

        std::vector<ParamGrad<T>> parameters() {
            std::vector<ParamGrad<T>> params;
            for (auto& layer : layers_) {
                auto lp = layer->parameters();
                params.insert(params.end(), lp.begin(), lp.end());
            }
            return params;
        }

    private:
        std::vector<std::unique_ptr<TransformerEncoderLayer<T>>> layers_;
    };

    // =================================================================
    // Transformer Decoder (N 層スタック)
    // =================================================================

    template<typename T>
    class TransformerDecoder {
    public:
        TransformerDecoder(size_t numLayers, size_t dModel, size_t numHeads,
                            size_t dFF, size_t seqLen, unsigned seed = 42)
        {
            for (size_t i = 0; i < numLayers; ++i)
                layers_.push_back(std::make_unique<TransformerDecoderLayer<T>>(
                    dModel, numHeads, dFF, seqLen, seed + static_cast<unsigned>(i * 10)));
        }

        void setEncoderOutput(const Matrix<T>& encOut) {
            for (auto& layer : layers_) layer->setEncoderOutput(encOut);
        }

        Matrix<T> forward(const Matrix<T>& input) {
            Matrix<T> x = input;
            for (auto& layer : layers_) x = layer->forward(x);
            return x;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) {
            Matrix<T> grad = gradOutput;
            for (int i = static_cast<int>(layers_.size()) - 1; i >= 0; --i)
                grad = layers_[i]->backward(grad);
            return grad;
        }

        std::vector<ParamGrad<T>> parameters() {
            std::vector<ParamGrad<T>> params;
            for (auto& layer : layers_) {
                auto lp = layer->parameters();
                params.insert(params.end(), lp.begin(), lp.end());
            }
            return params;
        }

    private:
        std::vector<std::unique_ptr<TransformerDecoderLayer<T>>> layers_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_TRANSFORMER_HPP
