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

int main(int argc, char *argv[])
{
	enum Command {Nop, Create, Reset, Train, Forward};
	Command command = Nop;
	
	if (argc > 1)
	{
		char cmd_token = *argv[1];
		if (cmd_token == 'c')
		{
			command = Command::Create;
		}
		else if (cmd_token == 'r')
		{
			command = Command::Reset;
		}
		else if (cmd_token == 't')
		{
			command = Command::Train;
		}
		else if (cmd_token == 'f')
		{
			command = Command::Forward;
		}
	}
	
	if (command == Command::Create)
	{
		arma::mat batch = RnnTextHelper::createTraining("moby_ch1.txt", RnnTextHelper::SEQ_LENGTH);
		batch.save("poet2.training.dat", arma::arma_ascii);
		return 0;
	}
	
	// Load project
	RnnStack *stack = reinterpret_cast<RnnStack*>(StackCreator::fromFile(".", "poet_2v_5s"));
	
	// Load weights
	stack->loadWeights(".");
	
	// Load training
	stack->loadTrainingBatch(".");
	arma::mat t_vc = stack->trainingBatch();

	if (command == Command::Reset)
	{
		for (int i=0; i < stack->numLayers(); i++)
		{
			stack->getLayer(i)->weightsInit(0.01);
		}
		stack->saveWeights(".");
	}
	
	if (command == Command::Train)
	{
		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 (command == Command::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;
	}
	printf("\nEnd of program\n");
	return 0;
}