// Network.cpp
// David R. Morrison

#include "network.hpp"

// Constructs a new network with the specified number of inputs and outputs;
// Each internal layer has a certain number of cell blocks, each with a given size
Network::Network(int in, int out, vector<int> blocks, vector<int> blockSize)
{
    inputVectorSize = in;
    outputVectorSize = out;

    // Initialize the random number generator
    initRand();
    assert(blocks.size() == blockSize.size());

    // All the layers except the output layer
    layers.resize(blocks.size());

    // The first layer takes in the inputs
    layers[0] = new CLayer(inputVectorSize, blocks[0], blockSize[0]);

    // All the others take the outputs of the previous layer
    for (int i = 1; i < layers.size(); ++i)
	layers[i] = new CLayer(blocks[i - 1] * blockSize[i - 1], blocks[i], blockSize[i]);

    // Output layer is just some linear neurons.
    outputLayer = new NLayer(blocks[blocks.size() - 1] * blockSize[blocks.size() - 1], 
	    outputVectorSize);
}

// Destroys the layers
Network::~Network()
{
    for (int i = 0; i < layers.size(); ++i)
	delete layers[i];

    delete outputLayer;
}

double Network::train(vector< vector< vector<double> > >& inputs,
	vector< vector< vector<double> > >& targets, double eta, int trace)
{
    assert(inputs.size() == targets.size());
    double mse = 0.0;

    for (int i = 0; i < inputs.size(); ++i)
    {
	for (int j = 0; j < inputVectorSize; ++j)
	{
	    vector<double> inp = layers[0]->getInputVector(inputs[i][j]);
	    cout << "Time Step " << j << ":" << endl;
	    cout << "Input: [ ";
	    for (int k = 0; k < inp.size(); ++k)
		cout << inp[k] << " ";
	    cout << "]" << endl;
	    double sampleError = 0.0;
	    vector<double> outputs = forwardPass(inputs[i][j]);
	    cout << "Network output: [ ";
	    for (int k = 0; k < outputs.size(); ++k)
		cout << outputs[k] << " ";
	    cout << "]" << endl << "Targets: [";
	    assert(outputs.size() == targets[i][j].size());
	    for (int k = 0; k < targets[i][j].size(); ++k)
	    {
		cout << targets[i][j][k] << " ";
		sampleError += (outputs[k] - targets[i][j][k]) * (outputs[k] - targets[i][j][k]);
	    }
	    cout << "]" << endl << endl;
	    sampleError /= outputs.size();
	    mse += sampleError;
	    backwardPass(targets[i][j]);
	    updateWeights(eta);
	    if (trace)
		cout << *this << endl;
	}
	clear();
    }

    mse /= inputs.size();

    return mse;
}

double Network::test(vector< vector< vector<double> > >& inputs,
	vector< vector< vector<double> > >& targets)
{
    assert(inputs.size() == targets.size());
    double mse = 0.0;

    clear();
    for (int i = 0; i < inputs.size(); ++i)
    {
	for (int j = 0; j < inputVectorSize; ++j)
	{            vector<double> inp = layers[0]->getInputVector(inputs[i][j]);
            cout << "Time Step " << j << ":" << endl;
            cout << "Input: [ ";
            for (int k = 0; k < inp.size(); ++k)
                cout << inp[k] << " ";
            cout << "]" << endl;
            double sampleError = 0.0;
            vector<double> outputs = forwardPass(inputs[i][j]);
            cout << "Network output: [ ";
            for (int k = 0; k < outputs.size(); ++k)
                cout << outputs[k] << " ";
            cout << "]" << endl << "Targets: [";
            assert(outputs.size() == targets[i][j].size());
            for (int k = 0; k < targets[i][j].size(); ++k)
            {
                cout << targets[i][j][k] << " ";
                sampleError += (outputs[k] - targets[i][j][k]) * (outputs[k] - targets[i][j][k]);
            }
            cout << "]" << endl << endl;
            sampleError /= outputs.size();
            mse += sampleError;
	}
    }

    return mse;
}

// Takes in some input values at a given time step, and runs them through the network
vector<double> Network::forwardPass(vector<double> inputs)
{
    // Calculate the output from the first layer
    vector<double> interm = layers[0]->calcLayerOutput(inputs);

    // If there are more layers, feed the output from the previous layer into it
    for (int i = 1; i < layers.size(); ++i)
	interm = layers[i]->calcLayerOutput(interm);

    // Feed the output from the final inside layer into the output layer, and return
    vector<double> output = outputLayer->calcLayerOutput(interm);

    return output;
}

// Calculate error signals
void Network::backwardPass(vector<double> targets)
{
    assert(targets.size() == outputLayer->size());

    // Calculate the signals for the output layer; this is relatively easy, 
    // since it just subtracts the actual value from the target value, and the
    // output layer is linear.
    vector<double> outputSignals = outputLayer->calcSignals(targets); 

    // Get a vector that contains the weights from the previous layer
    // to the output layer.  This is stored as
    // outputWeights[cell][outputUnit]
    vector< vector<double> > outputWeights = outputLayer->getInputWeights();
    // Run backwards through the network transferring error signals back.
    for (int i = layers.size() - 1; i >= 0; --i)
    {
	outputSignals = layers[i]->calcSignals(outputSignals, outputWeights);
	outputWeights.clear();
	outputWeights = layers[i]->getInputWeights();
    }
}

// Update weights throughout the network, with a learning rate of eta
void Network::updateWeights(double eta)
{
    for (int i = 0; i < layers.size(); ++i)
	layers[i]->updateWeights(eta);
    outputLayer->updateWeights(eta);
}

void Network::clear()
{
    for (int i = 0; i < layers.size(); ++i)
	layers[i]->clear();
}

ostream& Network::print(ostream& out)
{
    out << "Network has " << inputVectorSize << " input(s) and " << 
	outputVectorSize << " output(s)." << endl;
    layers.size() == 1 ? out << "There is 1 internal layer." : 
	out << "There are " << layers.size() << " internal layers.";
    out << endl << endl;
    for (int i = 0; i < layers.size(); ++i)
	out << "Layer " << i << ":" << endl << *layers[i];
    out << "Output Layer: " << endl << *outputLayer;

    out << endl;
    return out;
}

ostream& operator<<(ostream& out, Network& n)
{
    return n.print(out);
}