#pragma once

#include "./core/omp_config.h"

#include "./utils/vector.h"
#include "./utils/matrix.h"
#include "./utils/matcast.h"

#include "./numerics/matclip.h"
#include "./numerics/veclog.h"

#include "./Loss.h"


namespace neural_networks{

	template <typename Td, typename Ti>
	struct Loss_CategoricalCrossentrophy : Loss<Td, Ti>  {
		
			utils::Vector<Td> forward(const utils::Matrix<Td>& y_pred, const utils::Matrix<Ti>& y_true) override{
				
				utils::Vector<Td> correct_confidences(y_true.rows(), Td{0});
				utils::Matrix<Td> cast_y_true = utils::matcast<Td, Ti>(y_true);

				// Number of samles in a batch
				const uint64_t samples = y_true.rows();

				// Clip data to prevent dividning by 0
				// Clip both sides to not drag mean towards any value
				utils::Matrix<Td> y_pred_clipped = numerics::matclip(y_pred, Td{1e-7}, Td{1.0} - Td{1e-7});

				// Probabilities for taget values
				// only if categorical labes
				if (y_true.cols() == 1){
					for (uint64_t i = 0; i < y_true.rows(); ++i){
						const uint64_t idx = static_cast<uint64_t>(y_true(i, 0));
						correct_confidences[i] = y_pred_clipped(i, idx);
					}
				}else{ // Mask values - only for one-hot encoded labels
					correct_confidences = numerics::matdot_row(y_pred_clipped, cast_y_true);

				}
				// Losses			
				utils::Vector<Td> negative_log_likelihoods(samples, Td{0});
				for (uint64_t i = 0; i < samples; ++i){
				    negative_log_likelihoods[i] = -std::log(static_cast<Td>(correct_confidences[i]));
				}

				return negative_log_likelihoods;
			}
	};



} // end namespace neural_networks