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

// Layer.hpp — NN レイヤー抽象基底、Linear 層、Sequential コンテナ
//
// Matrix<T> をテンソルとして使用 (行=バッチ, 列=特徴量)。
// 各レイヤーは forward/backward を実装し、パラメータと勾配を保持する。

#ifndef CALX_ML_NN_LAYER_HPP
#define CALX_ML_NN_LAYER_HPP

#include <math/core/matrix.hpp>
#include <cassert>
#include <cmath>
#include <memory>
#include <random>
#include <utility>
#include <vector>

namespace calx::ml::nn {

    /// パラメータと勾配のペア
    template<typename T>
    struct ParamGrad {
        Matrix<T>* param;
        Matrix<T>* grad;
    };

    // =================================================================
    // Layer 基底クラス
    // =================================================================

    template<typename T>
    class Layer {
    public:
        virtual ~Layer() = default;

        /// 順伝播 (入力: batch × inFeatures → 出力: batch × outFeatures)
        virtual Matrix<T> forward(const Matrix<T>& input) = 0;

        /// 逆伝播 (勾配: batch × outFeatures → 入力勾配: batch × inFeatures)
        virtual Matrix<T> backward(const Matrix<T>& gradOutput) = 0;

        /// パラメータ/勾配のリスト (最適化対象)
        virtual std::vector<ParamGrad<T>> parameters() { return {}; }

        /// 訓練/推論モード切替
        virtual void setTraining(bool training) { training_ = training; }
        bool isTraining() const { return training_; }

    protected:
        bool training_ = true;
    };

    // =================================================================
    // Linear (全結合層)
    // =================================================================

    /// y = x * W^T + b  (入力 batch×in, 出力 batch×out)
    template<typename T>
    class Linear : public Layer<T> {
    public:
        /// @param inFeatures   入力特徴量数
        /// @param outFeatures  出力特徴量数
        /// @param useBias      バイアスを使うか
        /// @param seed         乱数シード
        Linear(size_t inFeatures, size_t outFeatures,
               bool useBias = true, unsigned seed = 42)
            : inFeatures_(inFeatures), outFeatures_(outFeatures), useBias_(useBias)
        {
            // Kaiming (He) 初期化
            T scale = std::sqrt(T{2} / static_cast<T>(inFeatures));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            weight_ = Matrix<T>(outFeatures, inFeatures);
            gradWeight_ = Matrix<T>(outFeatures, inFeatures, T{0});
            for (size_t i = 0; i < outFeatures; ++i)
                for (size_t j = 0; j < inFeatures; ++j)
                    weight_(i, j) = dist(rng);

            if (useBias_) {
                bias_ = Matrix<T>(1, outFeatures, T{0});
                gradBias_ = Matrix<T>(1, outFeatures, T{0});
            }
        }

        Matrix<T> forward(const Matrix<T>& input) override {
            input_ = input;  // backward 用にキャッシュ
            size_t batch = input.rows();

            // output = input * W^T
            Matrix<T> output(batch, outFeatures_);
            for (size_t b = 0; b < batch; ++b) {
                for (size_t o = 0; o < outFeatures_; ++o) {
                    T sum = T{0};
                    for (size_t i = 0; i < inFeatures_; ++i)
                        sum += input(b, i) * weight_(o, i);
                    output(b, o) = sum + (useBias_ ? bias_(0, o) : T{0});
                }
            }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = gradOutput.rows();

            // gradWeight = gradOutput^T * input
            for (size_t o = 0; o < outFeatures_; ++o) {
                for (size_t i = 0; i < inFeatures_; ++i) {
                    T sum = T{0};
                    for (size_t b = 0; b < batch; ++b)
                        sum += gradOutput(b, o) * input_(b, i);
                    gradWeight_(o, i) += sum;
                }
            }

            // gradBias = sum of gradOutput over batch
            if (useBias_) {
                for (size_t o = 0; o < outFeatures_; ++o) {
                    T sum = T{0};
                    for (size_t b = 0; b < batch; ++b)
                        sum += gradOutput(b, o);
                    gradBias_(0, o) += sum;
                }
            }

            // gradInput = gradOutput * W
            Matrix<T> gradInput(batch, inFeatures_);
            for (size_t b = 0; b < batch; ++b) {
                for (size_t i = 0; i < inFeatures_; ++i) {
                    T sum = T{0};
                    for (size_t o = 0; o < outFeatures_; ++o)
                        sum += gradOutput(b, o) * weight_(o, i);
                    gradInput(b, i) = sum;
                }
            }
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            std::vector<ParamGrad<T>> params;
            params.push_back({ &weight_, &gradWeight_ });
            if (useBias_)
                params.push_back({ &bias_, &gradBias_ });
            return params;
        }

        const Matrix<T>& weight() const { return weight_; }
        const Matrix<T>& bias() const { return bias_; }

    private:
        size_t inFeatures_, outFeatures_;
        bool useBias_;
        Matrix<T> weight_, bias_;
        Matrix<T> gradWeight_, gradBias_;
        Matrix<T> input_;  // forward 時のキャッシュ
    };

    // =================================================================
    // Dropout 層
    // =================================================================

    template<typename T>
    class Dropout : public Layer<T> {
    public:
        explicit Dropout(T rate = T{0.5}, unsigned seed = 42)
            : rate_(rate), rng_(seed) {}

        Matrix<T> forward(const Matrix<T>& input) override {
            if (!this->training_ || rate_ <= T{0}) return input;

            size_t rows = input.rows(), cols = input.cols();
            mask_ = Matrix<T>(rows, cols);
            std::bernoulli_distribution drop(static_cast<double>(rate_));
            T scale = T{1} / (T{1} - rate_);

            Matrix<T> output(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j) {
                    mask_(i, j) = drop(rng_) ? T{0} : scale;
                    output(i, j) = input(i, j) * mask_(i, j);
                }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            if (!this->training_ || rate_ <= T{0}) return gradOutput;

            size_t rows = gradOutput.rows(), cols = gradOutput.cols();
            Matrix<T> gradInput(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j)
                    gradInput(i, j) = gradOutput(i, j) * mask_(i, j);
            return gradInput;
        }

    private:
        T rate_;
        std::mt19937 rng_;
        Matrix<T> mask_;
    };

    // =================================================================
    // Sequential コンテナ
    // =================================================================

    template<typename T>
    class Sequential {
    public:
        Sequential() = default;

        /// レイヤーを追加
        template<typename LayerT, typename... Args>
        Sequential& add(Args&&... args) {
            layers_.push_back(std::make_unique<LayerT>(std::forward<Args>(args)...));
            return *this;
        }

        /// 既存レイヤーを移動追加
        Sequential& add(std::unique_ptr<Layer<T>> layer) {
            layers_.push_back(std::move(layer));
            return *this;
        }

        /// 順伝播
        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;
        }

        /// 訓練/推論モード切替
        void setTraining(bool training) {
            for (auto& layer : layers_)
                layer->setTraining(training);
        }

        size_t numLayers() const { return layers_.size(); }

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

} // namespace calx::ml::nn

#endif // CALX_ML_NN_LAYER_HPP
