#include "core/omp_config.h"

#include "utils/utils.h"
#include "numerics/numerics.h"
#include "decomp/decomp.h"

#include "modules/neural_networks/neural_networks.h"
#include "random/random.h"



//#include <iostream>
//#include <stdexcept>
//#include <chrono>





int main(int argc, char const *argv[])
{   

    uint64_t number_of_classes = 2;
    uint64_t number_of_samples = 1000;
    uint64_t number_of_epochs = 10000;

    utils::Mf X;
    utils::Mf X_test;
    utils::Matrix<float> y;
    utils::Matrix<float> y_test;
    float data_loss;
    float regularization_loss;
    float loss;
    float accuracy;

    utils::Vector<uint64_t> class_targets;
    utils::Matrix<float> predictions;


    // Create dataset
    //neural_networks::create_spital_data<float, int64_t>(number_of_samples, number_of_classes, X, y);
    //neural_networks::create_vertical_data<float, int64_t>(number_of_samples, number_of_classes, X, y);
    neural_networks::create_sine_data(number_of_samples,        // samples
                                     2.0f* 3.1415f,                   // length
                                     X,                         
                                     y);



    
    // Create Dense layer with 2 input featues and 3 output values
    neural_networks::Dense_Layer<float> dense1(
                                            1, 64,  // input/output
                                            0.0f,  // weight L1
                                            5e-5f,  // weight L2
                                            0.0f,   // bias L1
                                            5e-5f    // bias L2
                                            );
    neural_networks::Activation_ReLU<float> activation1;
    //neural_networks::Dropout_Layer<float> dropout1(0.1);

    neural_networks::Dense_Layer<float> dense2(
                                            64, 64,  // input/output
                                            0.0f,  // weight L1
                                            5e-5f,  // weight L2
                                            0.0f,   // bias L1
                                            5e-5f    // bias L2
                                            );
    neural_networks::Activation_ReLU<float> activation2;


    neural_networks::Dense_Layer<float> dense3(
                                            64, 1,  // input/output
                                            0.0f,  // weight L1
                                            5e-5f,  // weight L2
                                            0.0f,   // bias L1
                                            5e-5f    // bias L2
                                            );
    neural_networks::Activation_Linear<float> activation3;


    neural_networks::Loss_MeanSquaredError<float> loss_function;
    neural_networks::Optimizer_Adam<float> optimizer(
                                                    0.001,   // Learning-rate
                                                    1e-3,   // Learning-rate decay
                                                    1e-7,   // epsilons
                                                    0.9,    // beta 1 
                                                    0.999   // beta 2
                                                    );
    

    /* Accuracy precision for accuracy calculation
    # There are no really accuracy factor for regression problem,
    # but we can simulate/approximate it. We'll calculate it by checking
    # how many values have a difference to their ground truth equivalent
    # less than given precision
    # We'll calculate this precision as a fraction of standard deviation
    # of al the ground truth values */
    // accuracy_precision = np.std(y) / 250
    float accuracy_precision = numerics::standard_deviation(y)/ 250.0f;


    // Train in loop
    for (uint64_t epoch = 0; epoch < number_of_epochs+1; ++epoch){

        // Perform a forward pass of our training data through this layer
        dense1.forward(X);
        activation1.forward(dense1.outputs);
        //dropout1.forward(activation1.outputs);

        dense2.forward(activation1.outputs);
        activation2.forward(dense2.outputs);

        dense3.forward(activation2.outputs);
        activation3.forward(dense3.outputs);

        // Perform a foard pass through the activation/loss function
        // takes the output of the second dense layer here and returns loss
        data_loss = loss_function.calculate(activation3.outputs, y);

        // Calculate regularization penalty
        regularization_loss = loss_function.regularization_loss(dense1) + 
                            loss_function.regularization_loss(dense2);

        loss = data_loss + regularization_loss;

        predictions = activation3.outputs;
        accuracy = numerics::mean(numerics::less( numerics::abs( numerics::sub(predictions, y)), accuracy_precision));
        //accuracy = numerics::mean(numerics::equal_elementwise_serial(predictions, utils::veccast<float, int64_t>(y.get_col(0))));


        if (!(epoch%100) && epoch != 0){

            std::cout << "epoch: " << epoch;
            std::cout << ", acc: " << accuracy;
            std::cout << ", loss: " << loss;
            std::cout << ", data_loss: " << data_loss;
            std::cout << ", regularization_loss: " << regularization_loss;
            std::cout << ", lr: " << optimizer.current_learning_rate;
            std::cout << std::endl;

        }

        // Backward pass
        loss_function.backward(activation3.outputs, y);

        activation3.backward(loss_function.dinputs);
        dense3.backward(activation3.dinputs);

        activation2.backward(dense3.dinputs);
        dense2.backward(activation2.dinputs);

        activation1.backward(dense2.dinputs);
        dense1.backward(activation1.dinputs);

        // Update weights and biases
        optimizer.pre_update_params();
        optimizer.update_params(dense1);
        optimizer.update_params(dense2);
        optimizer.update_params(dense3);
        optimizer.post_update_params();

    }
/*
std::cout << "X, y, pred:" << std::endl;

for (uint64_t i = 0; i < X.rows(); ++i) {
    std::cout << X(i, 0)
              << ", "
              << y(i, 0)
              << ", "
              << activation3.outputs(i, 0)
              << std::endl;
}

    // Validate the model

    // Create dataset
    neural_networks::create_spital_data<float, int64_t>(100, number_of_classes, X_test, y_test);

    // Perform a forward pass of our training data through this layer
    dense1.forward(X_test);
    activation1.forward(dense1.outputs);
    //dropout1.forward(activation1.outputs);

    dense2.forward(activation1.outputs);
    activation2.forward(dense2.outputs);

    dense3.forward(activation2.outputs);
    activation3.forward(dense3.outputs);

    // Perform a foard pass through the activation/loss function
    // takes the output of the second dense layer here and returns loss
    data_loss = loss_activation.calculate(activation3.outputs, y_test);

    // Calculate regularization penalty
    regularization_loss = loss_activation.regularization_loss(dense1) + 
                        loss_activation.regularization_loss(dense2) + 
                        loss_activation.regularization_loss(dense3);

    loss = data_loss + regularization_loss;


    // skal flyttes ned under loss functions. 
    predictions = activation3.outputs();
    predictions = numerics::mean(numerics::abs(numerics::sub(predictions, y)));
    std::cout << predictions << std::endl;

    // Calculate accuracy from output of activation2 and targets
    predictions = numerics::greater_than(activation3.outputs, 0.5f).get_col(0);

    accuracy = numerics::mean(numerics::equal_elementwise_serial(predictions, utils::veccast<float, int64_t>(y_test.get_col(0))));


    std::cout << "validation, acc: " << accuracy << ", loss: " << loss << std::endl;
*/
    return 0;
}