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

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

const char punctuation[] = {' ', '.', '!', '?', 0};
const int NUM_CODES = 1 + 26 + 10;
const int SEQ_LENGTH = 2;

int char_is(char c, const char *pTable)
{
	int found = 0;
	while(*pTable)
	{
		if (c == *pTable)
		{
			found = 1;
			break;
		}
		pTable++;
	}
	return found;
}

int ch2idx(char c)
{
	c = toupper(c);

	int index = 0;
	if (isalpha(c))
	{
		index = 1+ (int)(c-'A');
	}
	if (isdigit(c))
	{
		index = 1+26+(int)(c-'0');
	}
	return index;
}

char idx2ch(int index)
{
	char c = '?';
	if (index == 0)
	{
		c = ' ';
	}
	else if (index >= 1 and index < 27)
	{
		c = (char)(index + 'A' - 1);
	}
	else if (index >= 27 and index < 38)
	{
		c = (char)(index + '0' - 27);
	}
	return c;
}

arma::mat createTraining(const string &filename, size_t seq_len)
{
	FILE *pFile;

	pFile = fopen(filename.c_str(), "r");

	if (!pFile)
	{
		std::cout << "Could not open " << filename << "!" << std::endl;
		return 0;
	}

	int numTraining = 0;
	int numVisible = seq_len*NUM_CODES;
	
	arma::mat batch = arma::zeros(numTraining, numVisible);
	arma::mat pattern = arma::zeros(seq_len, NUM_CODES);
	
	for (int i=0; i < seq_len; i++)
	{
		pattern.row(i)[0] = 1;
	}
	
	int last_index = ch2idx(' ');
	while(!feof(pFile))
	{
		char c;
		int result = fread(&c, 1, 1, pFile);
		if (result < 0)
		{
			break;
		}
		
		int index = ch2idx(c);
		if (index == 0 and last_index == 0)
		{
			continue;
		}
		last_index = index;
		pattern = arma::shift(pattern, 1);
		arma::mat data = arma::zeros(1, NUM_CODES);
		data[index] = 1;
		pattern.row(0) = data;
		arma::mat pattern_vector = pattern.as_row();
		batch.insert_rows(batch.n_rows, pattern_vector);
	}
	fclose(pFile);
	
	return batch;
}

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

int main()
{
#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();

	Rbm::Params params = stack->getLayer(0)->params();
	
#if DO_TRAINING
	int numGibbs = params.numGibbs;
	params.numGibbs = 1;
	stack->setParams(params);
	RbmListener listener;
	
	for (int i=0; i < stack->getSeqLen(); i++)
	{
		stack->getLayer(i)->weightsInit(0.01,0);
	}
	
	// Phase 10
	params.numEpochs = 2000;
	stack->setParams(params);
	stack->train(t_vc, &listener);

	// Phase 11
	params.numEpochs = 10000;
	stack->setParams(params);
	stack->train(t_vc, &listener);
	
	// Phase 20
	params.numGibbs = numGibbs;
	stack->setParams(params);
	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
	
	params.gibbsDoSampleHidden = false;
	params.gibbsDoSampleVisible = false;
	stack->setParams(params);
	arma::mat state;
	
	for (int i=0; i < t_vc.n_rows; i++)
	{
		arma::mat 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 = idx2ch(next.index_max());
		putchar(c);
	}
	printf("\n");

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