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

// BackPropagation.hpp — 多層パーセプトロン (バックプロパゲーション)
//
// 移植元: lib++20 機械学習/ニューラルネット/バックプロパゲーション.h
// 改良点:
//   - calx Matrix<T>/Vector<T> ベース
//   - std::vector<int> でレイヤー構成指定
//   - 乱数を calx ではなく std::mt19937 で初期化
//   - API 英語化
//   - 活性化関数: Sigmoid, ReLU, Tanh, Softmax (層ごとに指定可能)
//   - オプティマイザ: SGD (momentum), Adam
//   - ミニバッチ学習 (trainBatch)
//   - ドロップアウト

#ifndef CALX_ML_NEURAL_BACKPROPAGATION_HPP
#define CALX_ML_NEURAL_BACKPROPAGATION_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <numeric>
#include <random>
#include <vector>

namespace calx::ml {

    /// 活性化関数の種別
    enum class Activation { Sigmoid, ReLU, Tanh, Softmax };

    /// オプティマイザの種別
    enum class Optimizer { SGD, Adam };

    /// 多層パーセプトロン
    template<typename T>
    class MLP {
    public:
        /// 旧 API 互換コンストラクタ (Sigmoid + SGD momentum)
        /// @param layers  各層のニューロン数 (入力層, 隠れ層..., 出力層)
        /// @param lr      学習率
        /// @param momentum モメンタム係数
        /// @param weightLimit 重みクリッピング閾値
        /// @param seed    乱数シード
        MLP(const std::vector<int>& layers, T lr = T{0.5}, T momentum = T{0.1},
            T weightLimit = T{1}, unsigned seed = 42)
            : MLP(layers, std::vector<Activation>(layers.size() - 1, Activation::Sigmoid),
                  Optimizer::SGD, lr, momentum, weightLimit, T{0}, seed) {}

        /// 拡張コンストラクタ
        /// @param layers       各層のニューロン数
        /// @param activations  各層 (入力層を除く) の活性化関数
        /// @param optimizer    オプティマイザ種別
        /// @param lr           学習率
        /// @param momentum     SGD モメンタム / 未使用 (Adam)
        /// @param weightLimit  重みクリッピング閾値 (0 = 無効)
        /// @param dropout      ドロップアウト率 (0 = 無効, 隠れ層のみ適用)
        /// @param seed         乱数シード
        MLP(const std::vector<int>& layers,
            const std::vector<Activation>& activations,
            Optimizer optimizer = Optimizer::SGD,
            T lr = T{0.01}, T momentum = T{0.9},
            T weightLimit = T{0}, T dropout = T{0},
            unsigned seed = 42)
            : numLayers_(static_cast<int>(layers.size()))
            , neurons_(layers)
            , activations_(activations)
            , optimizer_(optimizer)
            , lr_(lr)
            , momentum_(momentum)
            , limit_(weightLimit)
            , dropout_(dropout)
            , rng_(seed)
        {
            assert(numLayers_ >= 2);
            assert(static_cast<int>(activations_.size()) == numLayers_ - 1);

            // 出力・δ・pre-activation
            out_.resize(numLayers_);
            delta_.resize(numLayers_);
            preAct_.resize(numLayers_);
            dropMask_.resize(numLayers_);
            for (int i = 0; i < numLayers_; ++i) {
                out_[i].resize(neurons_[i], T{0});
                delta_[i].resize(neurons_[i], T{0});
                preAct_[i].resize(neurons_[i], T{0});
                dropMask_[i].resize(neurons_[i], T{1});
            }

            // 重み初期化 (He / Xavier)
            weights_.resize(numLayers_);
            for (int i = 1; i < numLayers_; ++i) {
                int rows = neurons_[i];
                int cols = neurons_[i - 1] + 1;  // +1 for bias
                weights_[i] = Matrix<T>(rows, cols, T{0});

                // He init for ReLU, Xavier for others
                T scale;
                if (activations_[i - 1] == Activation::ReLU)
                    scale = std::sqrt(T{2} / neurons_[i - 1]);
                else
                    scale = std::sqrt(T{1} / neurons_[i - 1]);
                std::normal_distribution<T> dist(T{0}, scale);

                for (int j = 0; j < rows; ++j)
                    for (int k = 0; k < cols; ++k)
                        weights_[i](j, k) = dist(rng_);
            }

            initOptimizer();
        }

        /// 順伝播
        void forward(const T* input) {
            training_ = false;
            forwardImpl(input);
        }

        void forward(const Vector<T>& input) {
            forward(input.data());
        }

        /// 1 サンプルで学習 (順伝播 + 逆伝播 + 重み更新)
        void train(const T* input, const T* target) {
            training_ = true;
            forwardImpl(input);
            backprop(target);
            updateWeights();
        }

        void train(const Vector<T>& input, const Vector<T>& target) {
            train(input.data(), target.data());
        }

        /// ミニバッチ学習
        /// @param X       入力データ (サンプル数 × 入力次元)
        /// @param Y       ターゲット (サンプル数 × 出力次元)
        /// @param epochs  エポック数
        /// @param batchSize ミニバッチサイズ (0 = 全バッチ)
        void trainBatch(const Matrix<T>& X, const Matrix<T>& Y,
                        size_t epochs = 1, size_t batchSize = 32)
        {
            size_t n = X.rows();
            if (batchSize == 0 || batchSize > n) batchSize = n;

            std::vector<size_t> indices(n);
            std::iota(indices.begin(), indices.end(), 0);

            for (size_t epoch = 0; epoch < epochs; ++epoch) {
                std::shuffle(indices.begin(), indices.end(), rng_);

                for (size_t start = 0; start < n; start += batchSize) {
                    size_t end = std::min(start + batchSize, n);
                    size_t bs = end - start;

                    // 勾配を累積
                    resetGradients();

                    for (size_t s = start; s < end; ++s) {
                        size_t idx = indices[s];
                        training_ = true;
                        forwardImpl(&X(idx, 0));
                        backprop(&Y(idx, 0));
                        accumulateGradients();
                    }

                    // 平均勾配で更新
                    T invBs = T{1} / static_cast<T>(bs);
                    for (int i = 1; i < numLayers_; ++i) {
                        int rows = neurons_[i];
                        int cols = neurons_[i - 1] + 1;
                        for (int j = 0; j < rows; ++j)
                            for (int k = 0; k < cols; ++k)
                                gradAccum_[i](j, k) *= invBs;
                    }
                    updateWeightsFromGrad();
                }
            }
        }

        /// 出力層の i 番目の出力
        T output(int i) const { return out_[numLayers_ - 1][i]; }

        /// 出力層全体
        const std::vector<T>& outputs() const { return out_[numLayers_ - 1]; }

        /// 平均二乗誤差
        T mse(const T* target) const {
            int outLayer = numLayers_ - 1;
            T sum = T{0};
            for (int i = 0; i < neurons_[outLayer]; ++i) {
                T diff = target[i] - out_[outLayer][i];
                sum += diff * diff;
            }
            return sum / neurons_[outLayer];
        }

        /// 学習率の設定
        void setLearningRate(T lr) { lr_ = lr; }

    private:
        // --- 活性化関数 ---
        T activate(Activation act, T x) const {
            switch (act) {
            case Activation::Sigmoid: return T{1} / (T{1} + std::exp(-x));
            case Activation::ReLU:    return x > T{0} ? x : T{0};
            case Activation::Tanh:    return std::tanh(x);
            default:                  return x; // Softmax は別処理
            }
        }

        T activateDerivFromOutput(Activation act, T out, T preact) const {
            switch (act) {
            case Activation::Sigmoid: return out * (T{1} - out);
            case Activation::ReLU:    return preact > T{0} ? T{1} : T{0};
            case Activation::Tanh:    return T{1} - out * out;
            case Activation::Softmax: return out * (T{1} - out); // 簡略版 (cross-entropy 前提)
            default:                  return T{1};
            }
        }

        void applySoftmax(std::vector<T>& v) const {
            T maxVal = *std::max_element(v.begin(), v.end());
            T sum = T{0};
            for (auto& x : v) { x = std::exp(x - maxVal); sum += x; }
            for (auto& x : v) x /= sum;
        }

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

        // --- 順伝播の実装 ---
        void forwardImpl(const T* input) {
            for (int i = 0; i < neurons_[0]; ++i)
                out_[0][i] = input[i];

            // ドロップアウトマスク生成
            if (training_ && dropout_ > T{0}) {
                std::bernoulli_distribution drop(static_cast<double>(dropout_));
                for (int i = 1; i < numLayers_ - 1; ++i) { // 隠れ層のみ
                    for (int j = 0; j < neurons_[i]; ++j)
                        dropMask_[i][j] = drop(rng_) ? T{0} : T{1} / (T{1} - dropout_);
                }
            }

            for (int i = 1; i < numLayers_; ++i) {
                Activation act = activations_[i - 1];

                for (int j = 0; j < neurons_[i]; ++j) {
                    T sum = T{0};
                    for (int k = 0; k < neurons_[i - 1]; ++k)
                        sum += out_[i - 1][k] * weights_[i](j, k);
                    sum += weights_[i](j, neurons_[i - 1]);  // bias
                    preAct_[i][j] = sum;

                    if (act != Activation::Softmax)
                        out_[i][j] = activate(act, sum);
                    else
                        out_[i][j] = sum; // Softmax は後で一括処理
                }

                if (act == Activation::Softmax)
                    applySoftmax(out_[i]);

                // ドロップアウト適用 (隠れ層、訓練時のみ)
                if (training_ && dropout_ > T{0} && i < numLayers_ - 1) {
                    for (int j = 0; j < neurons_[i]; ++j)
                        out_[i][j] *= dropMask_[i][j];
                }
            }
        }

        // --- 逆伝播 ---
        void backprop(const T* target) {
            int outLayer = numLayers_ - 1;
            Activation outAct = activations_[outLayer - 1];

            // 出力層の δ
            if (outAct == Activation::Softmax) {
                // Softmax + cross-entropy: δ = target - out (勾配上昇方向に統一)
                for (int i = 0; i < neurons_[outLayer]; ++i)
                    delta_[outLayer][i] = target[i] - out_[outLayer][i];
            } else {
                for (int i = 0; i < neurons_[outLayer]; ++i) {
                    T error = target[i] - out_[outLayer][i];
                    delta_[outLayer][i] = activateDerivFromOutput(
                        outAct, out_[outLayer][i], preAct_[outLayer][i]) * error;
                }
            }

            // 隠れ層の δ
            for (int i = outLayer - 1; i > 0; --i) {
                Activation act = activations_[i - 1];
                for (int j = 0; j < neurons_[i]; ++j) {
                    T sum = T{0};
                    for (int k = 0; k < neurons_[i + 1]; ++k)
                        sum += delta_[i + 1][k] * weights_[i + 1](k, j);
                    delta_[i][j] = activateDerivFromOutput(
                        act, out_[i][j], preAct_[i][j]) * sum;

                    // ドロップアウトマスク
                    if (dropout_ > T{0} && i < numLayers_ - 1)
                        delta_[i][j] *= dropMask_[i][j];
                }
            }
        }

        // --- オプティマイザ ---
        void initOptimizer() {
            if (optimizer_ == Optimizer::SGD) {
                prevDelta_.resize(numLayers_);
                for (int i = 1; i < numLayers_; ++i)
                    prevDelta_[i] = Matrix<T>(neurons_[i], neurons_[i - 1] + 1, T{0});
            } else { // Adam
                adamM_.resize(numLayers_);
                adamV_.resize(numLayers_);
                for (int i = 1; i < numLayers_; ++i) {
                    int rows = neurons_[i], cols = neurons_[i - 1] + 1;
                    adamM_[i] = Matrix<T>(rows, cols, T{0});
                    adamV_[i] = Matrix<T>(rows, cols, T{0});
                }
                adamT_ = 0;
            }

            gradAccum_.resize(numLayers_);
            for (int i = 1; i < numLayers_; ++i)
                gradAccum_[i] = Matrix<T>(neurons_[i], neurons_[i - 1] + 1, T{0});
        }

        // 単一サンプル用: 即時重み更新
        void updateWeights() {
            if (optimizer_ == Optimizer::SGD) {
                // モメンタム
                if (momentum_ != T{0}) {
                    for (int i = 1; i < numLayers_; ++i)
                        for (int j = 0; j < neurons_[i]; ++j)
                            for (int k = 0; k <= neurons_[i - 1]; ++k)
                                weights_[i](j, k) += momentum_ * prevDelta_[i](j, k);
                }

                for (int i = 1; i < numLayers_; ++i) {
                    for (int j = 0; j < neurons_[i]; ++j) {
                        for (int k = 0; k < neurons_[i - 1]; ++k) {
                            T dw = lr_ * delta_[i][j] * out_[i - 1][k];
                            prevDelta_[i](j, k) = dw;
                            weights_[i](j, k) += dw;
                            clipWeight(weights_[i](j, k));
                        }
                        T dw = lr_ * delta_[i][j];
                        prevDelta_[i](j, neurons_[i - 1]) = dw;
                        weights_[i](j, neurons_[i - 1]) += dw;
                        clipWeight(weights_[i](j, neurons_[i - 1]));
                    }
                }
            } else { // Adam — 単一サンプル
                resetGradients();
                accumulateGradients();
                updateWeightsFromGrad();
            }
        }

        void resetGradients() {
            for (int i = 1; i < numLayers_; ++i) {
                int rows = neurons_[i], cols = neurons_[i - 1] + 1;
                for (int j = 0; j < rows; ++j)
                    for (int k = 0; k < cols; ++k)
                        gradAccum_[i](j, k) = T{0};
            }
        }

        void accumulateGradients() {
            for (int i = 1; i < numLayers_; ++i) {
                for (int j = 0; j < neurons_[i]; ++j) {
                    for (int k = 0; k < neurons_[i - 1]; ++k)
                        gradAccum_[i](j, k) += delta_[i][j] * out_[i - 1][k];
                    gradAccum_[i](j, neurons_[i - 1]) += delta_[i][j]; // bias
                }
            }
        }

        // 勾配累積済みの状態から重み更新
        void updateWeightsFromGrad() {
            if (optimizer_ == Optimizer::SGD) {
                for (int i = 1; i < numLayers_; ++i) {
                    int cols = neurons_[i - 1] + 1;
                    for (int j = 0; j < neurons_[i]; ++j) {
                        for (int k = 0; k < cols; ++k) {
                            T dw = lr_ * gradAccum_[i](j, k);
                            if (momentum_ != T{0}) {
                                dw += momentum_ * prevDelta_[i](j, k);
                                prevDelta_[i](j, k) = dw;
                            }
                            weights_[i](j, k) += dw;
                            clipWeight(weights_[i](j, k));
                        }
                    }
                }
            } else { // Adam
                ++adamT_;
                constexpr T beta1 = T{0.9}, beta2 = T{0.999}, eps = T{1e-8};
                T bc1 = T{1} - std::pow(beta1, static_cast<T>(adamT_));
                T bc2 = T{1} - std::pow(beta2, static_cast<T>(adamT_));

                for (int i = 1; i < numLayers_; ++i) {
                    int cols = neurons_[i - 1] + 1;
                    for (int j = 0; j < neurons_[i]; ++j) {
                        for (int k = 0; k < cols; ++k) {
                            T g = gradAccum_[i](j, k);
                            adamM_[i](j, k) = beta1 * adamM_[i](j, k) + (T{1} - beta1) * g;
                            adamV_[i](j, k) = beta2 * adamV_[i](j, k) + (T{1} - beta2) * g * g;
                            T mHat = adamM_[i](j, k) / bc1;
                            T vHat = adamV_[i](j, k) / bc2;
                            weights_[i](j, k) += lr_ * mHat / (std::sqrt(vHat) + eps);
                            clipWeight(weights_[i](j, k));
                        }
                    }
                }
            }
        }

        void clipWeight(T& w) const {
            if (limit_ > T{0} && std::abs(w) > limit_)
                w *= T{0.999};
        }

        int numLayers_;
        std::vector<int> neurons_;
        std::vector<Activation> activations_;
        Optimizer optimizer_;
        T lr_, momentum_, limit_, dropout_;
        std::mt19937 rng_;
        bool training_ = false;

        std::vector<std::vector<T>> out_;
        std::vector<std::vector<T>> delta_;
        std::vector<std::vector<T>> preAct_;
        std::vector<std::vector<T>> dropMask_;
        std::vector<Matrix<T>> weights_;

        // SGD momentum
        std::vector<Matrix<T>> prevDelta_;

        // Adam
        std::vector<Matrix<T>> adamM_, adamV_;
        int adamT_ = 0;

        // ミニバッチ勾配累積
        std::vector<Matrix<T>> gradAccum_;
    };

} // namespace calx::ml

#endif // CALX_ML_NEURAL_BACKPROPAGATION_HPP
