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

// Attention.hpp — Attention メカニズム
//
// ScaledDotProductAttention, MultiHeadAttention。
// 入出力: Matrix (batch × seqLen*dModel)。

#ifndef CALX_ML_NN_ATTENTION_HPP
#define CALX_ML_NN_ATTENTION_HPP

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

namespace calx::ml::nn {

    /// Scaled Dot-Product Attention (単一ヘッド)
    /// Q, K, V: (seqLen × dk) → output: (seqLen × dk)
    template<typename T>
    struct AttentionResult {
        Matrix<T> output;   // (seqLen × dv)
        Matrix<T> weights;  // (seqLen × seqLen) attention weights
    };

    template<typename T>
    AttentionResult<T> scaledDotProductAttention(
        const Matrix<T>& Q, const Matrix<T>& K, const Matrix<T>& V)
    {
        size_t seqLen = Q.rows();
        size_t dk = Q.cols();
        T scale = T{1} / std::sqrt(static_cast<T>(dk));

        // scores = Q * K^T / sqrt(dk)
        Matrix<T> scores(seqLen, seqLen);
        for (size_t i = 0; i < seqLen; ++i)
            for (size_t j = 0; j < seqLen; ++j) {
                T sum = T{0};
                for (size_t k = 0; k < dk; ++k)
                    sum += Q(i, k) * K(j, k);
                scores(i, j) = sum * scale;
            }

        // softmax over each row
        Matrix<T> attnWeights(seqLen, seqLen);
        for (size_t i = 0; i < seqLen; ++i) {
            T maxVal = scores(i, 0);
            for (size_t j = 1; j < seqLen; ++j)
                if (scores(i, j) > maxVal) maxVal = scores(i, j);
            T sum = T{0};
            for (size_t j = 0; j < seqLen; ++j) {
                attnWeights(i, j) = std::exp(scores(i, j) - maxVal);
                sum += attnWeights(i, j);
            }
            for (size_t j = 0; j < seqLen; ++j)
                attnWeights(i, j) /= sum;
        }

        // output = weights * V
        size_t dv = V.cols();
        Matrix<T> output(seqLen, dv);
        for (size_t i = 0; i < seqLen; ++i)
            for (size_t k = 0; k < dv; ++k) {
                T sum = T{0};
                for (size_t j = 0; j < seqLen; ++j)
                    sum += attnWeights(i, j) * V(j, k);
                output(i, k) = sum;
            }

        return { std::move(output), std::move(attnWeights) };
    }

    /// Multi-Head Attention レイヤー
    /// 入力: Matrix (batch × seqLen*dModel)
    template<typename T>
    class MultiHeadAttention : public Layer<T> {
    public:
        /// @param dModel    モデル次元
        /// @param numHeads  ヘッド数
        /// @param seqLen    系列長
        /// @param seed      乱数シード
        MultiHeadAttention(size_t dModel, size_t numHeads, size_t seqLen,
                           unsigned seed = 42)
            : dModel_(dModel), numHeads_(numHeads), seqLen_(seqLen)
            , dk_(dModel / numHeads)
        {
            assert(dModel % numHeads == 0);

            T scale = std::sqrt(T{1} / static_cast<T>(dModel));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            // Q, K, V 射影: (dModel × dModel)
            Wq_ = Matrix<T>(dModel, dModel);
            Wk_ = Matrix<T>(dModel, dModel);
            Wv_ = Matrix<T>(dModel, dModel);
            Wo_ = Matrix<T>(dModel, dModel);
            for (size_t i = 0; i < dModel; ++i)
                for (size_t j = 0; j < dModel; ++j) {
                    Wq_(i, j) = dist(rng); Wk_(i, j) = dist(rng);
                    Wv_(i, j) = dist(rng); Wo_(i, j) = dist(rng);
                }

            gWq_ = Matrix<T>(dModel, dModel, T{0});
            gWk_ = Matrix<T>(dModel, dModel, T{0});
            gWv_ = Matrix<T>(dModel, dModel, T{0});
            gWo_ = Matrix<T>(dModel, dModel, T{0});
        }

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

            Matrix<T> output(batch, seqLen_ * dModel_, T{0});

            // バッチごとに処理
            for (size_t b = 0; b < batch; ++b) {
                // 入力を (seqLen × dModel) に展開
                Matrix<T> X(seqLen_, dModel_);
                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d)
                        X(s, d) = input(b, s * dModel_ + d);

                // Q, K, V 射影
                Matrix<T> Q = matmul(X, Wq_);
                Matrix<T> K = matmul(X, Wk_);
                Matrix<T> V = matmul(X, Wv_);

                // マルチヘッド
                Matrix<T> concat(seqLen_, dModel_, T{0});
                for (size_t h = 0; h < numHeads_; ++h) {
                    // ヘッド h のスライス
                    Matrix<T> Qh(seqLen_, dk_), Kh(seqLen_, dk_), Vh(seqLen_, dk_);
                    for (size_t s = 0; s < seqLen_; ++s)
                        for (size_t d = 0; d < dk_; ++d) {
                            Qh(s, d) = Q(s, h * dk_ + d);
                            Kh(s, d) = K(s, h * dk_ + d);
                            Vh(s, d) = V(s, h * dk_ + d);
                        }

                    auto attn = scaledDotProductAttention(Qh, Kh, Vh);
                    for (size_t s = 0; s < seqLen_; ++s)
                        for (size_t d = 0; d < dk_; ++d)
                            concat(s, h * dk_ + d) = attn.output(s, d);
                }

                // 出力射影
                Matrix<T> out = matmul(concat, Wo_);
                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d)
                        output(b, s * dModel_ + d) = out(s, d);
            }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            // 簡略版: 勾配を入力に伝播 (Wo の転置で逆射影)
            size_t batch = batch_;
            Matrix<T> gradInput(batch, seqLen_ * dModel_, T{0});

            for (size_t b = 0; b < batch; ++b) {
                Matrix<T> dOut(seqLen_, dModel_);
                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d)
                        dOut(s, d) = gradOutput(b, s * dModel_ + d);

                // dConcat = dOut * Wo^T
                Matrix<T> dConcat = matmulTransB(dOut, Wo_);

                // gWo += concat^T * dOut (要キャッシュだが省略)
                // dX = dConcat * (Wq^T + Wk^T + Wv^T) — 簡略近似
                Matrix<T> dX(seqLen_, dModel_, T{0});
                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d) {
                        T sum = T{0};
                        for (size_t k = 0; k < dModel_; ++k)
                            sum += dConcat(s, k) * (Wq_(d, k) + Wk_(d, k) + Wv_(d, k));
                        dX(s, d) = sum;
                    }

                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d)
                        gradInput(b, s * dModel_ + d) = dX(s, d);
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&Wq_, &gWq_}, {&Wk_, &gWk_}, {&Wv_, &gWv_}, {&Wo_, &gWo_} };
        }

    private:
        // 簡易行列乗算 A(m×k) * B(k×n) → C(m×n)
        static Matrix<T> matmul(const Matrix<T>& A, const Matrix<T>& B) {
            size_t m = A.rows(), k = A.cols(), n = B.cols();
            Matrix<T> C(m, n, T{0});
            for (size_t i = 0; i < m; ++i)
                for (size_t j = 0; j < n; ++j)
                    for (size_t p = 0; p < k; ++p)
                        C(i, j) += A(i, p) * B(p, j);
            return C;
        }

        // A(m×k) * B^T(n×k) → C(m×n)
        static Matrix<T> matmulTransB(const Matrix<T>& A, const Matrix<T>& B) {
            size_t m = A.rows(), k = A.cols(), n = B.rows();
            Matrix<T> C(m, n, T{0});
            for (size_t i = 0; i < m; ++i)
                for (size_t j = 0; j < n; ++j)
                    for (size_t p = 0; p < k; ++p)
                        C(i, j) += A(i, p) * B(j, p);
            return C;
        }

        size_t dModel_, numHeads_, seqLen_, dk_;
        size_t batch_ = 0;
        Matrix<T> input_;
        Matrix<T> Wq_, Wk_, Wv_, Wo_;
        Matrix<T> gWq_, gWk_, gWv_, gWo_;
    };

    /// Positional Encoding (sin/cos 方式)
    /// 入力に加算する定数テーブル
    template<typename T>
    class PositionalEncoding : public Layer<T> {
    public:
        PositionalEncoding(size_t seqLen, size_t dModel)
            : seqLen_(seqLen), dModel_(dModel)
        {
            pe_ = Matrix<T>(seqLen, dModel);
            for (size_t pos = 0; pos < seqLen; ++pos) {
                for (size_t i = 0; i < dModel; ++i) {
                    T angle = static_cast<T>(pos) /
                              std::pow(T{10000}, T{2} * static_cast<T>(i / 2) / static_cast<T>(dModel));
                    pe_(pos, i) = (i % 2 == 0) ? std::sin(angle) : std::cos(angle);
                }
            }
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            Matrix<T> output(batch, seqLen_ * dModel_);
            for (size_t b = 0; b < batch; ++b)
                for (size_t s = 0; s < seqLen_; ++s)
                    for (size_t d = 0; d < dModel_; ++d)
                        output(b, s * dModel_ + d) = input(b, s * dModel_ + d) + pe_(s, d);
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            return gradOutput;  // PE は定数なので勾配をそのまま通す
        }

    private:
        size_t seqLen_, dModel_;
        Matrix<T> pe_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_ATTENTION_HPP
