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

static const char *DEFAULT_START_WORD = "CHAP";

class RbmListener : public Rbm::IListener
{
	public:
	RbmListener(AStack &stack)
	: m_stack(stack)
	, m_last_progress(0)
	{
	}
	virtual ~RbmListener() {}
	
	bool onProgress(Rbm *pRbm, const Rbm::Status &status)
	{
		if (status.progress == 0)
		{
			m_last_progress = -10;
		}
		if ((status.progress - m_last_progress) == 10)
		{
			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;
			m_last_progress = status.progress;
			forward(DEFAULT_START_WORD, 100);
		}
		return true;
	}
	
	void forward(std::string const &start, size_t len)
	{
		RnnStack *stack = reinterpret_cast<RnnStack*>(&m_stack);
		arma::mat state;
		arma::mat curr;
		arma::mat next;
		std::string curr_str;
		
		for (int i=0; i < start.size(); i++)
		{
			curr = Matutils::char2vec(start.at(i), RnnTextHelper::NUM_CODES);
			curr_str.append(1, Matutils::vec2char(curr));
			arma::mat r = stack->step_forward(state, curr);
			next = stack->to_next(r);
			stack->clamp_one_hot(next);
		}
		cout << "Start: " << curr_str << std::endl;

		for (int i=stack->getSeqLen(); i < len; i++)
		{
			curr = next;
			curr_str.append(1, Matutils::vec2char(curr));
			arma::mat r = stack->step_forward(state, curr);
			next = stack->to_next(r);
			stack->clamp_one_hot(next);
			curr = stack->to_curr(r);
		}
		cout << "Curr: " << curr_str << std::endl;

	}
	AStack &m_stack;
	int m_last_progress;
};

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;
		}
	}
	
	// Load project
	RnnStack *stack = reinterpret_cast<RnnStack*>(StackCreator::fromFile(".", "poet_2v_5s"));
	
	// Load weights
	stack->loadWeights(".");
	
	RbmListener listener(*stack);
	
	if (command == Command::Reset)
	{
		stack->weightsInit(0.01);
		stack->saveWeights(".");
	}
	
	if (command == Command::Train)
	{
		arma::mat batch;
		if (argv[2] != nullptr)
		{
			batch = RnnTextHelper::createTraining(argv[2], RnnTextHelper::SEQ_LENGTH);
		}
		else
		{
			batch = RnnTextHelper::createTraining("batch1.txt", RnnTextHelper::SEQ_LENGTH);
		}
		
		// Phase 10
		stack->train(batch, &listener);

		stack->saveWeights(".");
#if DO_SAVE_PROJECT_AFTER_TRAINING
		stack->save(".");		
#endif
	}
	
	if (command == Command::Forward)
	{
		std::string start(DEFAULT_START_WORD);
		
		if (argv[2] != nullptr)
		{
			start = std::string(argv[2]);
		}
		listener.forward(start, 100);
	}
	printf("\nEnd of program\n");
	return 0;
}