#include <iostream>
#include <fstream>
#include <streambuf>
#include <cstdio>
#include <cmath>
#include <string>
#include <armadillo>
#include <jsoncpp/json/json.h>
#include "Rbm.hpp"
#include "Layer.hpp"
#include "Stack.hpp"

using namespace std;
class RbmListener : public Rbm::IListener
{
	public:
	RbmListener() {}
	virtual ~RbmListener() {}
		
	bool onProgress(const Rbm::Status &status)
	{
		std::cout << "Progress :          " << 100*status.progress << " %" << std::endl;
		std::cout << "epoch :             " << status.epoch << std::endl;
		std::cout << "trainingSizeRemain: " << status.trainingSizeRemain << std::endl;
		std::cout << "error (per mini batch) = " << status.err << std::endl;
		std::cout << "error (total)          = " << status.err_total << std::endl;
		std::cout << "L1                     = " << status.L1 << std::endl;
		std::cout << "L2                     = " << status.L2 << std::endl;
		
		return true;
	}
};

arma::mat loadTraining(const string &filename)
{
	uint32_t numTraining = 0;
	uint32_t numVisible = 0;
	FILE *pFile;

	pFile = fopen(filename.c_str(), "r");

	if (!pFile)
	{
		std::cout << "Could not open " << filename << "!" << std::endl;
		return 0;
	}

	int result = fscanf(pFile, "%d\n", &numTraining);
	if (result < 0)
	{
		return 0;
	}
	result = fscanf(pFile, "%d\n", &numVisible);
	if (result < 0)
	{
		return 0;
	}
	arma::mat data = arma::zeros(numTraining, numVisible);

	uint32_t i, j;
	for (i=0; i < numTraining; i++)
	{
		for (j=0; j < numVisible; j++)
		{
			float v;
			int result = fscanf(pFile, "%f", &v);
			if (result > 0)
			{
				data(i, j) = v;
			}
		}
	}
	fclose(pFile);
	return data;
	
}

int main()
{
	printf("Hallo, Welt!\n");

	const string project("test");
	
	RbmListener statusDisplay;
	Stack stack(project);

	arma::mat batch = loadTraining(project + string(".training.dat"));

	size_t numTraining = batch.n_rows;
	printf("Loaded %d training samples\n", (int)numTraining);

#if 0
	int i = 0;
	Layer *lowerLayer = new Layer("Layer", i, 16, 16, 8);
	stack.addLayer(lowerLayer);
	
	for (++i; i < 4; i++)
	{
		Layer *layer = new Layer("Layer", i, lowerLayer->bh().n_elem, 1, lowerLayer->bh().n_elem >> 1);
		lowerLayer = layer;
		stack.addLayer(layer);
	}

	stack.getLayer(0)->params().learningRate = 0.04;
	stack.getLayer(0)->params().numEpochs = 1000;

	stack.getLayer(1)->params().learningRate = 0.03;
	stack.getLayer(1)->params().numEpochs = 500;

	stack.getLayer(2)->params().learningRate = 0.02;
	stack.getLayer(2)->params().numEpochs = 250;

	stack.getLayer(3)->params().learningRate = 0.01;
	stack.getLayer(3)->params().numEpochs = 125;

	// Save project
	stack.save();

	// Shake weights
	stack.weightsInit(0.01);

	// Save weights
	stack.saveWeights();

#else	
	// Load project
	stack.load();
	
	// Load weights
	stack.loadWeights();
#endif	
	
	// Train stack
	stack.train(batch, &statusDisplay);
	
	// Save weights
	stack.saveWeights();

	Layer *layer = stack.getLayer(0);
	arma::mat v = arma::randu(numTraining, layer->bv().n_elem);
	arma::mat h = layer->toHiddenProbs(v);
	arma::mat r = layer->toVisibleProbs(h);
	return 0;
}