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

// Activation.hpp — 活性化関数レイヤー
//
// ReLU, Sigmoid, Tanh, Softmax, GELU を Layer として実装。

#ifndef CALX_ML_NN_ACTIVATION_HPP
#define CALX_ML_NN_ACTIVATION_HPP

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

namespace calx::ml::nn {

    // =================================================================
    // ReLU
    // =================================================================

    template<typename T>
    class ReLU : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override {
            input_ = input;
            size_t rows = input.rows(), cols = input.cols();
            Matrix<T> output(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j)
                    output(i, j) = input(i, j) > T{0} ? input(i, j) : T{0};
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            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) = input_(i, j) > T{0} ? gradOutput(i, j) : T{0};
            return gradInput;
        }

    private:
        Matrix<T> input_;
    };

    // =================================================================
    // Sigmoid
    // =================================================================

    template<typename T>
    class Sigmoid : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override {
            size_t rows = input.rows(), cols = input.cols();
            output_ = Matrix<T>(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j)
                    output_(i, j) = T{1} / (T{1} + std::exp(-input(i, j)));
            return output_;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            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) {
                    T s = output_(i, j);
                    gradInput(i, j) = gradOutput(i, j) * s * (T{1} - s);
                }
            return gradInput;
        }

    private:
        Matrix<T> output_;
    };

    // =================================================================
    // Tanh
    // =================================================================

    template<typename T>
    class TanhLayer : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override {
            size_t rows = input.rows(), cols = input.cols();
            output_ = Matrix<T>(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j)
                    output_(i, j) = std::tanh(input(i, j));
            return output_;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            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) {
                    T t = output_(i, j);
                    gradInput(i, j) = gradOutput(i, j) * (T{1} - t * t);
                }
            return gradInput;
        }

    private:
        Matrix<T> output_;
    };

    // =================================================================
    // Softmax
    // =================================================================

    template<typename T>
    class Softmax : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override {
            size_t rows = input.rows(), cols = input.cols();
            output_ = Matrix<T>(rows, cols);
            for (size_t b = 0; b < rows; ++b) {
                T maxVal = input(b, 0);
                for (size_t j = 1; j < cols; ++j)
                    if (input(b, j) > maxVal) maxVal = input(b, j);
                T sum = T{0};
                for (size_t j = 0; j < cols; ++j) {
                    output_(b, j) = std::exp(input(b, j) - maxVal);
                    sum += output_(b, j);
                }
                for (size_t j = 0; j < cols; ++j)
                    output_(b, j) /= sum;
            }
            return output_;
        }

        /// Softmax の backward は cross-entropy loss と組み合わせて使う前提
        /// gradOutput は通常 (output - target) が入る
        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            // cross-entropy + softmax の組み合わせでは gradOutput をそのまま通す
            return gradOutput;
        }

    private:
        Matrix<T> output_;
    };

    // =================================================================
    // GELU (Gaussian Error Linear Unit)
    // =================================================================

    template<typename T>
    class GELU : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override {
            input_ = input;
            size_t rows = input.rows(), cols = input.cols();
            Matrix<T> output(rows, cols);
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j) {
                    T x = input(i, j);
                    // GELU(x) = x * Phi(x), Phi = 0.5*(1+erf(x/sqrt(2)))
                    output(i, j) = x * T{0.5} * (T{1} + std::erf(x / std::sqrt(T{2})));
                }
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t rows = gradOutput.rows(), cols = gradOutput.cols();
            Matrix<T> gradInput(rows, cols);
            constexpr T sqrt2pi_inv = T{0.3989422804014327};  // 1/sqrt(2*pi)
            for (size_t i = 0; i < rows; ++i)
                for (size_t j = 0; j < cols; ++j) {
                    T x = input_(i, j);
                    T phi = T{0.5} * (T{1} + std::erf(x / std::sqrt(T{2})));
                    T pdf = sqrt2pi_inv * std::exp(T{-0.5} * x * x);
                    gradInput(i, j) = gradOutput(i, j) * (phi + x * pdf);
                }
            return gradInput;
        }

    private:
        Matrix<T> input_;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_ACTIVATION_HPP
