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

// BatchNorm.hpp — バッチ正規化
//
// BatchNorm1D: 全結合層用 (batch × features)
// BatchNorm2D: 畳み込み層用 (batch × C*H*W, チャネルごとに正規化)

#ifndef CALX_ML_NN_BATCHNORM_HPP
#define CALX_ML_NN_BATCHNORM_HPP

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

namespace calx::ml::nn {

    /// BatchNorm1D (全結合層用)
    template<typename T>
    class BatchNorm1D : public Layer<T> {
    public:
        /// @param numFeatures  特徴量数
        /// @param eps          数値安定化項
        /// @param momentum     移動平均の減衰率
        explicit BatchNorm1D(size_t numFeatures, T eps = T{1e-5}, T momentum = T{0.1})
            : nf_(numFeatures), eps_(eps), mom_(momentum)
        {
            gamma_ = Matrix<T>(1, nf_, T{1});
            beta_ = Matrix<T>(1, nf_, T{0});
            gradGamma_ = Matrix<T>(1, nf_, T{0});
            gradBeta_ = Matrix<T>(1, nf_, T{0});
            runMean_ = Matrix<T>(1, nf_, T{0});
            runVar_ = Matrix<T>(1, nf_, T{1});
        }

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

            if (this->training_) {
                // バッチ統計量
                mean_.resize(nf_);
                var_.resize(nf_);
                for (size_t j = 0; j < nf_; ++j) {
                    T m = T{0};
                    for (size_t b = 0; b < batch; ++b) m += input(b, j);
                    m /= static_cast<T>(batch);
                    mean_[j] = m;
                    T v = T{0};
                    for (size_t b = 0; b < batch; ++b) {
                        T d = input(b, j) - m;
                        v += d * d;
                    }
                    v /= static_cast<T>(batch);
                    var_[j] = v;

                    // 移動平均更新
                    runMean_(0, j) = (T{1} - mom_) * runMean_(0, j) + mom_ * m;
                    runVar_(0, j) = (T{1} - mom_) * runVar_(0, j) + mom_ * v;
                }

                // 正規化
                xNorm_ = Matrix<T>(batch, nf_);
                for (size_t b = 0; b < batch; ++b)
                    for (size_t j = 0; j < nf_; ++j) {
                        xNorm_(b, j) = (input(b, j) - mean_[j]) / std::sqrt(var_[j] + eps_);
                        output(b, j) = gamma_(0, j) * xNorm_(b, j) + beta_(0, j);
                    }
            } else {
                // 推論時: 移動平均を使用
                for (size_t b = 0; b < batch; ++b)
                    for (size_t j = 0; j < nf_; ++j) {
                        T xn = (input(b, j) - runMean_(0, j)) / std::sqrt(runVar_(0, j) + eps_);
                        output(b, j) = gamma_(0, j) * xn + beta_(0, j);
                    }
            }
            return output;
        }

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

            for (size_t j = 0; j < nf_; ++j) {
                T invStd = T{1} / std::sqrt(var_[j] + eps_);

                // gradGamma, gradBeta
                T dGamma = T{0}, dBeta = T{0};
                for (size_t b = 0; b < batch; ++b) {
                    dGamma += gradOutput(b, j) * xNorm_(b, j);
                    dBeta += gradOutput(b, j);
                }
                gradGamma_(0, j) += dGamma;
                gradBeta_(0, j) += dBeta;

                // gradInput
                T dVar = T{0}, dMean = T{0};
                for (size_t b = 0; b < batch; ++b)
                    dVar += gradOutput(b, j) * gamma_(0, j) * xNorm_(b, j);
                dVar *= T{-0.5} * invStd * invStd;

                for (size_t b = 0; b < batch; ++b) {
                    T dxNorm = gradOutput(b, j) * gamma_(0, j);
                    gradInput(b, j) = dxNorm * invStd
                                    + dVar * T{2} * xNorm_(b, j) * std::sqrt(var_[j] + eps_) * invN;
                }
                for (size_t b = 0; b < batch; ++b) dMean += gradInput(b, j);
                dMean *= -invN;
                // 簡略化: dMean を加算
                for (size_t b = 0; b < batch; ++b)
                    gradInput(b, j) += dMean;
            }
            return gradInput;
        }

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

    private:
        size_t nf_;
        T eps_, mom_;
        Matrix<T> gamma_, beta_;
        Matrix<T> gradGamma_, gradBeta_;
        Matrix<T> runMean_, runVar_;
        std::vector<T> mean_, var_;
        Matrix<T> xNorm_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_BATCHNORM_HPP
