/**++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 *
 *                                              PANIC
 *                              Portable Algorithms and Numerics In C++
 *
 *                       Scientific computing from scratch, with feeling.
 *
 * Copyright (c) 2026 Michelle Bausager
 *
 * This file is part of PANIC.
 *
 * PANIC is free software licensed under the GNU General Public License v3.0 or later.
 * You may redistribute and/or modify it under the terms of the GPL.
 *
 * PANIC is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the LICENSE file for the full license text.
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 *
 *++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 *
 * Project Name:      PANIC
 * Module Name:       neural_network
 * File Name:         loss_categorical_crossentropy.hpp
 * Revision:          0.1.0
 * Date:              30-07-2026
 * Author:            Michelle Bausager
 *
 * Description:
 * Defines the base loss_categorical_crossentropy used in neural network
 *
 *++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
//---------------------------------------------------------------------------------------------------------------------------
// INCLUDE DESCRIPTION
//---------------------------------------------------------------------------------------------------------------------------
#include <neural_network/loss/loss_categorical_crossentropy.hpp>
#include <config/omp.hpp>
#include <config/types.hpp> // panic::uint_t, panic::int_t, and panic::real_t
#include <tensor/matrix.hpp> // panic::tensor::real_matrix (uint_matrix, int_matrix)
#include <tensor/vector.hpp>

#include <math/clip.hpp>
#include <math/log.hpp>
#include <math/mul.hpp>
#include <tensor/generators/one_hot.hpp>
#include <math/div.hpp>

#include <math/sum.hpp>

//---------------------------------------------------------------------------------------------------------------------------
// PRIVATE CONSTANTS
//---------------------------------------------------------------------------------------------------------------------------
/**
 * @brief Minimum number of element operations before using the OpenMP-enabled loop.
 *
 * Small vectors and matrices are kept serial because the overhead of starting
 * worker threads can be larger than the work itself.
 */
static const panic::types::uint_t loss_categorical_crossentropy_omp_min_size = 500;

static const panic::types::real_t clip_min = 1e-7;
static const panic::types::real_t clip_max = 1 - 1e-7;
static const panic::types::real_t neg = -1;

//---------------------------------------------------------------------------------------------------------------------------


namespace panic{
	namespace neural_network{

//--------------------------------------------------------------------------------------------------------------------------
// Function Name : panic::neural_network::loss_categorical_crossentropy::forward
//
// Description:
//      Default implementation for catecorical labels. Derived classes can override it.
//--------------------------------------------------------------------------------------------------------------------------
bool loss_categorical_crossentropy::forward(
    const panic::tensor::real_matrix& y_pred,
    const panic::tensor::uint_vector& y_true){

    // Number of samples in a batch
    const panic::types::uint_t samples = y_pred.rows();

    if (samples != y_true.size()){
        return false;
    }


    // clip data to prevent log by 0
    // Clip both sides to not drag the mean towards any value
    panic::tensor::real_matrix y_pred_cliped = panic::math::clip(y_pred, clip_min, clip_max);

    // Vector to hold the correct confidences
    panic::tensor::real_vector correct_confidences(samples);

    PANIC_OMP_PARALLEL_FOR_IF(samples > loss_categorical_crossentropy_omp_min_size)
    for (panic::types::uint_t i = 0; i < samples; ++i){
        const panic::types::uint_t idx = y_true[i];

        //if (idx >= y_pred.cols()){
        //    return false;
        //}

        correct_confidences[i] = y_pred_cliped(i, idx);
    }

    // Calculate losses
    panic::tensor::real_vector negative_log_likelihoos(samples);

    negative_log_likelihoos = panic::math::mul(panic::math::log(correct_confidences), neg);

    sample_losses = negative_log_likelihoos;


    return true;
}


//--------------------------------------------------------------------------------------------------------------------------
// Function Name : panic::neural_network::loss_categorical_crossentropy::forward
//
// Description:
//      Default one-hot encoded labels implementation. Derived classes can override it.
//--------------------------------------------------------------------------------------------------------------------------
bool loss_categorical_crossentropy::forward(
    const panic::tensor::real_matrix& y_pred,
    const panic::tensor::real_matrix& y_true){


    if (y_pred.rows() != y_true.rows() || y_pred.cols() != y_true.cols()){
        return false;
    }

    // Number of samples in a batch
    const panic::types::uint_t samples = y_pred.rows();


    // clip data to prevent log by 0
    // Clip both sides to not drag the mean towards any value
    panic::tensor::real_matrix y_pred_cliped = panic::math::clip(y_pred, clip_min, clip_max);

    // Vector to hold the correct confidences
    panic::tensor::real_vector correct_confidences(samples);

    
    correct_confidences = panic::math::sum_rowwise(panic::math::mul(y_pred_cliped, y_true));


    // Calculate losses
    panic::tensor::real_vector negative_log_likelihoos(samples);

    negative_log_likelihoos = panic::math::mul(panic::math::log(correct_confidences), neg);

    sample_losses = negative_log_likelihoos;

    return true;
}





//--------------------------------------------------------------------------------------------------------------------------
// Function Name : panic::neural_network::loss_categorical_crossentropy::backward
//
// Description:
//      Default implementation for catecorical labels. Derived classes can override it.
//--------------------------------------------------------------------------------------------------------------------------
bool loss_categorical_crossentropy::backward(
    const panic::tensor::real_matrix& dvalues,
    const panic::tensor::uint_vector& y_true){

    if (dvalues.rows() == 0 || dvalues.cols() == 0 || y_true.size() != dvalues.rows()){
        return false;
    }

    panic::tensor::real_matrix y_true_one_hot;

    // Transforms it to one_hot
    if (! panic::tensor::one_hot(dvalues.cols(), y_true, y_true_one_hot)){
        return false;
    }

    // Uses other backward overloaded function to handle the rest.
    return backward(dvalues, y_true_one_hot);
}




//--------------------------------------------------------------------------------------------------------------------------
// Function Name : panic::neural_network::loss_categorical_crossentropy::backward
//
// Description:
//      Default one-hot encoded labels implementation. Derived classes can override it.
//--------------------------------------------------------------------------------------------------------------------------
bool loss_categorical_crossentropy::backward(
    const panic::tensor::real_matrix& dvalues,
    const panic::tensor::real_matrix& y_true){

    // Number of samples
    const panic::types::uint_t samples = dvalues.rows();

    if (samples == 0 || dvalues.cols() == 0 || y_true.rows() != dvalues.rows() || y_true.cols() != dvalues.cols()){
        return false;
    }

    // dinputs = y_true / dvalues
    if (!panic::math::div(y_true, dvalues, dinputs)){
        return false;
    }

    // scale = -1 / samples
    const panic::types::real_t scale = -static_cast<panic::types::real_t>(1) / static_cast<panic::types::real_t>(samples);


    // dinputs *= scale
    if (!panic::math::mul(dinputs, scale, dinputs)){
        return false;
    }

    return true;
}


	} // namespace tensor
} // namespace panic