#pragma once

#include "core/omp_config.h"

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

#include "numerics/mean.h"
#include "numerics/abs.h"

namespace neural_networks{

	template <typename Td, typename Ti = Td>
	struct Loss{

		utils::Vector<Td> sample_losses;
		utils::Matrix<Td> dinputs;
		Td data_loss;
		Td regularization_losss;

		virtual utils::Vector<Td> forward(const utils::Matrix<Td>& output, const utils::Matrix<Ti>& y) = 0;
		virtual void backward(const utils::Matrix<Td>& dvalues, const utils::Matrix<Ti>& y) = 0;

		Td calculate(const utils::Matrix<Td>& output, const utils::Matrix<Ti>& y){
			
			// Calculate sample losses
			sample_losses = forward(output, y);

			// Calculate mean loss
			data_loss = numerics::mean(sample_losses);

			return data_loss;

		}

		template <typename Layer>
		Td regularization_loss(const Layer& layer){
			// 0 by default
			regularization_losss = 0;
			
			// L1 regularization - weights
			// calculate only when factor greater than 0
			if (layer.weight_regularizer_l1){
				regularization_losss += layer.weight_regularizer_l1 * numerics::sum(numerics::abs(layer.weights));
			}

			// L2 regularization - weights
			if (layer.weight_regularizer_l2){
				regularization_losss += layer.weight_regularizer_l2 * numerics::sum(numerics::mul(layer.weights,layer.weights)); // elementwise!
			}

			// L1 regularization - biases
			// calculate only when factor greater than 0
			if (layer.bias_regularizer_l1){
				regularization_losss += layer.bias_regularizer_l1 * layer.biases.abs().sum();
			}
			// L2 regularization - biases
			if (layer.bias_regularizer_l2){
				regularization_losss += layer.bias_regularizer_l2 * layer.biases.multiply(layer.biases).sum();
			}

			return regularization_losss;
		}


	};

} // end namespace neural_networks