#include "layer.h"
#include "matrices.h"
#include <vector>

class Network {
    public:
    Matrix input;
    float (*activation)(float) = Layer::Sigmoid; // Activation function is sigmoid by default

    std::vector<Layer> hidden_layers;
    Layer output_layer;

    inline void Feed(Matrix);
    inline Matrix GetOutput();

    inline void Forward();

    inline void BackPropagate(Matrix);

    // Constructors
    // Input size, Array of hidden sizes, Output size
    Network(int, std::vector<int>, int);
};

void Network::BackPropagate(Matrix target){
    // Calculate derivative of loss in respect to A (dca) for output layer
    // loss = (A - Y)^2
    // derivative = 2(A - Y)
    Matrix loss = this->output_layer.activated_output.Substract(&target);
    loss = loss.Hadamard(&loss);
    // loss.Print("Loss");
    Matrix dca = this->output_layer.activated_output.Substract(&target);
    dca = dca.Multiply(2.0F);
    // dca.Print("DCA");

    this->output_layer.BackPropagate(this->hidden_layers[this->hidden_layers.size() - 1].activated_output, dca, &Layer::SigmoidPrime);
}

Network::Network(int input_size, std::vector<int> hidden_sizes, int output_size){
    this->input = Matrix(input_size, 1);

    this->hidden_layers.push_back(Layer(input_size, hidden_sizes[0]));
    for(int i = 1; i < hidden_sizes.size(); i++){
        // For every hidden layer, create a layer of specified size
        this->hidden_layers.push_back(Layer(hidden_sizes[i-1], hidden_sizes[i]));
    }
    this->output_layer = Layer(hidden_sizes[hidden_sizes.size() - 1], output_size);
}

Matrix Network::GetOutput(){
    return this->output_layer.activated_output;
}

void Network::Feed(Matrix a){
    this->input = a;
}

void Network::Forward(){
    // Feeding first layer
    this->hidden_layers[0].Feed(this->input);
    this->hidden_layers[0].Forward();
    for(int i = 1; i < this->hidden_layers.size(); i++){
        // Feeding A(L-1) and forwarding
        this->hidden_layers[i].Feed(this->hidden_layers[i - 1].activated_output);
        this->hidden_layers[i].Forward();
    }
    this->output_layer.Feed(this->hidden_layers[this->hidden_layers.size() - 1].activated_output);
    this->output_layer.Forward();
}