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

using namespace std;
using namespace arma;

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;
	}
};


#define CREATE_TRAINING		0
#define DO_TRAINING			0
#define DO_FORWARD			1

int main()
{
	arma::mat sm = arma::zeros(8, 4);
	
	for (int i=0; i < sm.n_rows; i++)
	{
		sm = arma::shift(sm, 1, 0);
		//sm(0,0) = i;
		sm.row(0) = i*arma::eye(1,4);
		sm.print("sm");
	}
	
#if CREATE_TRAINING	
	arma::mat batch = createTraining("moby_ch1.txt", SEQ_LENGTH);
	batch.save("poet2.training.dat", arma::arma_ascii);
	return 0;
#endif
	
	// Load project
	RnnStack *stack = reinterpret_cast<RnnStack*>(StackCreator::fromFile(".", "poet2"));
	
	// Load weights
	stack->loadWeights(".");
	
	// Load training
	stack->loadTrainingBatch(".");
	arma::mat t_vc = stack->trainingBatch();

#if DO_TRAINING
	RbmListener listener;
	
	// Phase 10
	stack->train(t_vc, &listener);

	stack->saveWeights(".");
	stack->save(".");		

	printf("Curr\n");
	for (int i=0; i < t_vc.n_rows; i++)
	{
		arma::mat curr = stack->to_curr(t_vc.row(i));
		char c = idx2ch(curr.index_max());
		putchar(c);
	}
	printf("\n");
	printf("Next\n");
	for (int i=0; i < t_vc.n_rows; i++)
	{
		arma::mat next = stack->to_next(t_vc.row(i));
		char c = idx2ch(next.index_max());
		putchar(c);
	}
	printf("\n");

#endif
	
#if DO_FORWARD
	arma::mat state;
	arma::mat curr;
	for (int i=0; i < t_vc.n_rows; i++)
	{
		curr = stack->to_curr(t_vc.row(i));
		arma::mat r = stack->step_forward(state, curr);
		arma::mat next = stack->to_next(r);
		char c = RnnTextHelper::idx2ch(next.index_max());
		putchar(c);
	}
	printf("\n");
	for (int i=0; i < 40; i++)
	{
		curr = stack->to_next(stack->step_forward(state, curr));
		char c = RnnTextHelper::idx2ch(curr.index_max());
		putchar(c);
	}
	printf("\n");
#endif

	printf("\nEnd of program\n");
	return 0;
}