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

// Loss.hpp — 損失関数
//
// MSE, CrossEntropy, BinaryCrossEntropy。
// 各関数はスカラー損失値を返し、勾配 Matrix を出力引数で返す。

#ifndef CALX_ML_NN_LOSS_HPP
#define CALX_ML_NN_LOSS_HPP

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

namespace calx::ml::nn {

    /// MSE (Mean Squared Error) 損失
    /// @param output   予測 (batch × out)
    /// @param target   正解 (batch × out)
    /// @param grad     [出力] 勾配 (batch × out)
    /// @return         スカラー損失値
    template<typename T>
    T mseLoss(const Matrix<T>& output, const Matrix<T>& target, Matrix<T>& grad) {
        assert(output.rows() == target.rows() && output.cols() == target.cols());
        size_t n = output.rows(), d = output.cols();
        grad = Matrix<T>(n, d);

        T loss = T{0};
        T invND = T{1} / static_cast<T>(n * d);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < d; ++j) {
                T diff = output(i, j) - target(i, j);
                loss += diff * diff;
                grad(i, j) = T{2} * diff * invND;
            }
        return loss * invND;
    }

    /// Cross-Entropy 損失 (Softmax 出力前提)
    /// @param output   Softmax 後の確率 (batch × classes)
    /// @param target   One-hot ターゲット (batch × classes)
    /// @param grad     [出力] 勾配 = output - target (Softmax + CE の簡略化)
    /// @return         スカラー損失値
    template<typename T>
    T crossEntropyLoss(const Matrix<T>& output, const Matrix<T>& target, Matrix<T>& grad) {
        assert(output.rows() == target.rows() && output.cols() == target.cols());
        size_t n = output.rows(), d = output.cols();
        grad = Matrix<T>(n, d);

        T loss = T{0};
        constexpr T eps = T{1e-12};
        T invN = T{1} / static_cast<T>(n);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < d; ++j) {
                T p = std::max(output(i, j), eps);
                loss -= target(i, j) * std::log(p);
                // Softmax + CE の勾配: (output - target) / batch
                grad(i, j) = (output(i, j) - target(i, j)) * invN;
            }
        return loss * invN;
    }

    /// Binary Cross-Entropy 損失 (Sigmoid 出力前提)
    /// @param output   Sigmoid 後の確率 (batch × 1)
    /// @param target   0/1 ターゲット (batch × 1)
    /// @param grad     [出力] 勾配
    /// @return         スカラー損失値
    template<typename T>
    T binaryCrossEntropyLoss(const Matrix<T>& output, const Matrix<T>& target, Matrix<T>& grad) {
        assert(output.rows() == target.rows() && output.cols() == target.cols());
        size_t n = output.rows(), d = output.cols();
        grad = Matrix<T>(n, d);

        T loss = T{0};
        constexpr T eps = T{1e-12};
        T invND = T{1} / static_cast<T>(n * d);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < d; ++j) {
                T p = std::clamp(output(i, j), eps, T{1} - eps);
                T t = target(i, j);
                loss -= t * std::log(p) + (T{1} - t) * std::log(T{1} - p);
                // Sigmoid + BCE の勾配: (output - target) / (output * (1 - output)) * sigmoid_deriv
                // = (output - target) / batch
                grad(i, j) = (output(i, j) - target(i, j)) * invND;
            }
        return loss * invND;
    }

} // namespace calx::ml::nn

#endif // CALX_ML_NN_LOSS_HPP
