#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 "Stack.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 numCodes = 1 + 26 + 10;

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)
{
	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 = 2*numCodes;
	
	arma::mat batch = arma::zeros(numTraining, numVisible);
	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;
		}
		arma::mat data = arma::zeros(1, numVisible);
		
		data[index+numCodes] = 1;
		data[last_index] = 1;
		last_index = index;
		batch.insert_rows(batch.n_rows, data);
	}
	fclose(pFile);
	
	return batch;
}

arma::mat to_curr(arma::mat v)
{
	return v.cols(0, numCodes-1);
}

arma::mat to_next(arma::mat v)
{
	return v.cols(numCodes, 2*numCodes-1);
}

#define CREATE_TRAINING		0
int main()
{
#if CREATE_TRAINING	
	arma::mat batch = createTraining("moby_ch1.txt");
	batch.save("poet.training.dat", arma::arma_ascii);
#endif
	
	Stack stack(".", "poet2");
	
	// Load project
	stack.load();
	
	// Load weights
	stack.loadWeights();
	
	// Load training
	stack.loadTrainingBatch();

	Layer *layer = stack.getLayer(0);
	int numTraining = stack.trainingBatch().n_rows;
	
	arma::mat v = stack.trainingBatch();
	
	arma::mat h = layer->toHiddenProbs(v);
	arma::mat r = layer->toVisibleProbs(h);
	
	r = layer->upDownPass(v);
	layer->params().gibbsDoSampleHidden = false;
	layer->params().gibbsDoSampleVisible = false;
	
	printf("\nStimulus: Training\n");
	for (int i=0; i < numTraining; i++)
	{
		arma::mat curr = to_curr(v.row(i));
		char c = idx2ch(curr.index_max());
		putchar(c);
	}

	h = layer->toHiddenProbs(v.row(0));
	printf("\nStimulus: Next char and current h\n");
	for (int i=1; i < numTraining+1; i++)
	{
		layer->gibbs_hv(h, r);
		arma::mat curr = to_curr(r);
		arma::mat next = to_next(r);
		v = arma::join_rows(next, arma::zeros(1, next.n_cols), h);
		
		layer->gibbs_vh(v, h);
		
		char c = idx2ch(curr.index_max());
		putchar(c);
	}

	h = layer->toHiddenProbs(v.row(0));
	printf("\nStimulus: Current h\n");
	for (int i=1; i < numTraining+1; i++)
	{
		layer->gibbs_hv(h, r);
		arma::mat curr = to_curr(r);
		arma::mat next = to_next(r);
		v = arma::join_rows(arma::zeros(1, next.n_cols), arma::zeros(1, next.n_cols), h);
		layer->gibbs_vh(v, h);
		
		char c = idx2ch(curr.index_max());
		putchar(c);
	}

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