#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			1
#define DO_FORWARD			1

int main()
{
#if CREATE_TRAINING	
	arma::mat batch = RnnTextHelper::createTraining("moby_ch1.txt", RnnTextHelper::SEQ_LENGTH);
	batch.save("poet2.training.dat", arma::arma_ascii);
	return 0;
#endif
	
	// Load project
	RnnStack *stack = reinterpret_cast<RnnStack*>(StackCreator::fromFile(".", "poet_2v_5s"));
	
	// Load weights
	bool weight_loaded = stack->loadWeights(".");
	
	// Load training
	stack->loadTrainingBatch(".");
	arma::mat t_vc = stack->trainingBatch();

	bool do_training = not weight_loaded;
#if DO_TRAINING
	do_training = true;
#endif
			
	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 = RnnTextHelper::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 = RnnTextHelper::idx2ch(next.index_max());
			putchar(c);
		}
		printf("\n");
	}
	
#if DO_FORWARD
	arma::mat state;
	arma::mat curr;
	arma::mat next;
	std::string curr_str;
	std::string next_str;
	for (int i=0; i < stack->getSeqLen(); i++)
	{
		curr = stack->to_curr(t_vc.row(i));
		curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
		arma::mat r = stack->step_forward(state, curr);
		next = stack->to_next(r, true);
		next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
	}
	cout << "Curr: " << curr_str << std::endl;
	cout << "Next: " << next_str << std::endl;
	
	for (int i=stack->getSeqLen(); i < t_vc.n_rows; i++)
	{
		curr = next;
		curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
		arma::mat r = stack->step_forward(state, curr);
		next = stack->to_next(r, true);
		curr = stack->to_curr(r);
		next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
	}
	cout << "Curr: " << curr_str << std::endl;
	cout << "Next: " << next_str << std::endl;
#endif

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