#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 "DeepStack.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 = 5;

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

struct Rnn
{
	Rnn(Layer *layer)
	: nV(layer->numVisible() - layer->context().n_cols)
	, nH(layer->numHidden())
	, nVx(layer->numVisibleX())
	, nVy(layer->numVisibleY())
	, nC(layer->context().n_cols)
	{
		
	}
	size_t nV;
	size_t nH;
	size_t nVx;
	size_t nVy;
	size_t nC;
	
	arma::mat vcVec_to_vMat(arma::mat const &vcVec)
	{
		arma::mat vVec = vcVec.submat(0, 0, 0, nV-1);
		arma::mat vMat = arma::reshape(vVec, nVx, nVy);
		return vMat;
	}
	
	arma::mat curr_vec(arma::mat const &vcVec)
	{
		arma::mat vMat = vcVec_to_vMat(vcVec);
		return vMat.col(0);
	}
	
	arma::mat vcVec_next_step(arma::mat const &vCurr, arma::mat const &vcVec_last, arma::mat const &h)
	{
		arma::mat vMat_last = arma::shift(vcVec_to_vMat(vcVec_last), 1, 1);
		vMat_last.col(0) = arma::zeros(nVx, 1);
		vMat_last.col(1) = vCurr;
		
		arma::mat vVec = vMat_last.as_row();
		
		arma::mat vcVec_next = arma::join_rows(vVec, h);
		return vcVec_next;
	}
	
};

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

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

#define CREATE_TRAINING		0
int main()
{
#if CREATE_TRAINING	
	arma::mat batch = createTraining("moby_ch1.txt", SEQ_LENGTH);
	batch.save("poet.training.dat", arma::arma_ascii);
	return 0;
#endif
	
	DeepStack stack(".", "poet5");
	
	// 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 t = stack.trainingBatch();
	
	arma::mat h = layer->toHiddenProbs(t);
	arma::mat r = layer->toVisibleProbs(h);
	
	layer->params().gibbsDoSampleHidden = false;
	layer->params().gibbsDoSampleVisible = false;
	
	Rnn rnn(layer);
	
	printf("\nStimulus: Training\n");
	h = layer->toHiddenProbs(t);
	r = t;
	for (int i=0; i < numTraining; i++)
	{
		arma::mat r = t.row(i);
		layer->gibbs_vh(r, h);
		layer->gibbs_hv(h, r);
		arma::mat curr = rnn.curr_vec(r);
		char c = idx2ch(curr.index_max());
		putchar(c);
	}
	printf("\n");

	arma::mat v = t.row(0);
	printf("Stimulus: Next char and current h\n");
	for (int i=1; i < numTraining+1; i++)
	{
		arma::mat r = v;
		layer->gibbs_vh(r, h);
		arma::mat curr = rnn.curr_vec(r);
		v = rnn.vcVec_next_step(curr, r, h);
		
		layer->gibbs_vh(v, h);

		char c = idx2ch(curr.index_max());
		putchar(c);
	}
	printf("\n");
	
	v = t.row(0);
	h = layer->toHiddenProbs(v);
	printf("Stimulus: Current h\n");
	for (int i=1; i < numTraining+1; i++)
	{
		arma::mat r = v;
		layer->gibbs_hv(h, r);
		arma::mat curr = rnn.curr_vec(r);
		v = rnn.vcVec_next_step(curr, arma::zeros(1, r.n_cols), h);

		layer->gibbs_vh(v, h);
		
		char c = idx2ch(curr.index_max());
		putchar(c);
	}

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