#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 "DeepStack.hpp"

using namespace std;
class RbmListener : public Rbm::IListener
{
	public:
	RbmListener() {}
	virtual ~RbmListener() {}
	
	bool onProgress(Rbm *pRbm, const Rbm::Status &status)
	{
		std::cout << "Progress :          " << status.progress << " %" << 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;
	}
};

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

	const string project("context99");
	
	DeepStack stack(".", project);

#if CREATE_TEST
	stack.addTraining(stack.trainingBatch().row(1));
	printf("There are %d training samples\n", (int)stack.trainingBatch().n_rows);
	
	stack.delTraining(0);
	printf("There are %d training samples\n", (int)stack.trainingBatch().n_rows);

	const int numLayers = 4;
	int i = 0;
	Layer *lowerLayer = new Layer("Layer", i, 16, 16, 8);
	stack.addLayer(lowerLayer);
	
	for (++i; i < numLayers; i++)
	{
		Layer *layer = new Layer("Layer", i, lowerLayer->bh().n_elem, 1, lowerLayer->bh().n_elem >> 1);
		layer->params().learningRate = lowerLayer->params().learningRate/2;
		layer->params().numEpochs = lowerLayer->params().numEpochs/2;
		lowerLayer = layer;
		stack.addLayer(layer);
	}

	// Save project
	stack.save();

	// Shake weights
	stack.weightsInit(0.01);

	// Save weights
	stack.saveWeights();

#else	
	// Load project
	stack.load();
	
	// Load weights
	stack.loadWeights();
	
	// Load training
	stack.loadTrainingBatch();

#endif	
	
#if TRAIN_TEST
	RbmListener statusDisplay;

	// Train stack
	stack.train(&statusDisplay);
	
	// Save weights
	stack.saveWeights();
#endif
	Layer *layer = stack.getLayer(0);
	arma::mat v = stack.trainingBatch();
	v.print("t");
	layer->calcContextBatch(v);
	arma::mat h = layer->toHiddenProbs(v);
	arma::mat r = layer->toVisibleProbs(h);
	
	v.print("v1");
	r = layer->upDownPass(v);
	r.print("v2");
	
	arma::mat err = Rbm::rms_error(r-v);
	err.print("err");
	
	
	return 0;
}