//
// Created by martin on 7/15/18.
//

#include <vector>

#include "ErrorFunctions.h"

namespace lib4neuro {

    size_t ErrorFunction::get_dimension() {
        return this->dimension;
    }

    MSE::MSE(NeuralNetwork *net, DataSet *ds) {
        this->net = net;
        this->ds = ds;
        this->dimension = net->get_n_weights() + net->get_n_biases();
    }

    double MSE::eval(std::vector<double> *weights) {
        size_t dim_out = this->ds->get_output_dim();
//    unsigned int dim_in = this->ds->get_input_dim();
        size_t n_elements = this->ds->get_n_elements();
        double error = 0.0, val;

        std::vector<std::pair<std::vector<double>, std::vector<double>>> *data = this->ds->get_data();

//    //TODO instead use something smarter
//    this->net->copy_weights(weights);

        std::vector<double> output(dim_out);

        for (auto el: *data) {  // Iterate through every element in the test set

            this->net->eval_single(el.first, output,
                                   weights);  // Compute the net output and store it into 'output' variable


//        printf("errors: ");
            for (size_t j = 0; j < dim_out; ++j) {  // Compute difference for every element of the output vector

                val = output[j] - el.second[j];
                error += val * val;

//            printf("%f, ", val * val);
            }
//        printf("\n");

        }

//    printf("n_elements: %d\n", n_elements);
        return error / n_elements;
    }

    void MSE::calculate_error_gradient(std::vector<double> &params, std::vector<double> &grad, double alpha) {

        size_t dim_out = this->ds->get_output_dim();
        size_t n_elements = this->ds->get_n_elements();

        std::vector<std::pair<std::vector<double>, std::vector<double>>> *data = this->ds->get_data();

        std::vector<double> error_derivative(dim_out);


        for (auto el: *data) {  // Iterate through every element in the test set

            this->net->eval_single(el.first, error_derivative,
                                   &params);  // Compute the net output and store it into 'output' variable

            for (size_t j = 0; j < dim_out; ++j) {
                error_derivative[j] = 2.0 * (error_derivative[j] - el.second[j]); //real - expected result
            }

            this->net->add_to_gradient_single(el.first, error_derivative, alpha / n_elements, grad);
        }
    }

    std::vector<double> *MSE::get_parameters() {
        std::vector<double> *output = new std::vector<double>(this->net->get_n_weights() + this->net->get_n_biases());

        size_t i = 0;

        for (auto el: *this->net->get_parameter_ptr_weights()) {
            output->at(i) = el;
            ++i;
        }

        for (auto el: *this->net->get_parameter_ptr_biases()) {
            output->at(i) = el;
            ++i;
        }

        return output;
    }

    ErrorSum::ErrorSum() {
        this->summand = nullptr;
        this->summand_coefficient = nullptr;
        this->dimension = 0;
    }

    ErrorSum::~ErrorSum() {
        if (this->summand) {
            delete this->summand;
        }
        if (this->summand_coefficient) {
            delete this->summand_coefficient;
        }
    }

    double ErrorSum::eval(std::vector<double> *weights) {
        double output = 0.0;
        ErrorFunction *ef = nullptr;

        for (unsigned int i = 0; i < this->summand->size(); ++i) {
            ef = this->summand->at(i);

            if (ef) {
                output += ef->eval(weights) * this->summand_coefficient->at(i);
            }
        }

        return output;
    }

    void ErrorSum::calculate_error_gradient(std::vector<double> &params, std::vector<double> &grad, double alpha) {

        ErrorFunction *ef = nullptr;
        for (size_t i = 0; i < this->summand->size(); ++i) {
            ef = this->summand->at(i);

            if (ef) {
                ef->calculate_error_gradient(params, grad, this->summand_coefficient->at(i) * alpha);
            }
        }
    }

    void ErrorSum::add_error_function(ErrorFunction *F, double alpha) {
        if (!this->summand) {
            this->summand = new std::vector<ErrorFunction *>(0);
        }
        this->summand->push_back(F);

        if (!this->summand_coefficient) {
            this->summand_coefficient = new std::vector<double>(0);
        }
        this->summand_coefficient->push_back(alpha);

        if (F) {
            if (F->get_dimension() > this->dimension) {
                this->dimension = F->get_dimension();
            }
        }
    }

    size_t ErrorSum::get_dimension() {
//    if(!this->dimension) {
//        size_t max = 0;
//        for(auto e : *this->summand) {
//            if(e->get_dimension() > max) {
//                max = e->get_dimension();
//            }
//        };
//
//        this->dimension = max;
//    }
        return this->dimension;
    }

    std::vector<double> *ErrorSum::get_parameters() {
        return this->summand->at(0)->get_parameters();
    }

}