/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */

/* 
 * File:   DeepStack.cpp
 * Author: jens
 * 
 * Created on 25. Oktober 2019, 18:26
 */
#include <cassert>
#include "DeepStack.hpp"

using namespace std;

DeepStack::DeepStack(const std::string &dir, const std::string &name, StackType type)
: AStack(dir, type, name)
{
}

DeepStack::DeepStack(const DeepStack& orig)
: AStack(orig)
{
}

DeepStack::~DeepStack()
{ 
}

void DeepStack::train(const arma::mat& batch, Rbm::IListener* pListener)
{
	arma::mat thisBatch = batch;
	Layer *pLayer = getLayer(0);
	while (pLayer)
	{
		thisBatch = trainingBatchFrom(pLayer->id(), thisBatch);
		pLayer->train(thisBatch, pListener);
		pLayer = pLayer->next;
	}
}

void DeepStack::train(size_t layerId, const arma::mat& batch, Rbm::IListener* pListener)
{
	arma::mat thisBatch = batch;
	Layer *pLayer = getLayer(0);
	while (pLayer)
	{
		if (pLayer->id() == layerId)
		{
			break;
		}
		thisBatch = pLayer->toHiddenProbs(thisBatch);
		pLayer = pLayer->next;
	}
	pLayer->train(thisBatch, pListener);
}

arma::mat& DeepStack::trainingBatch()
{
	return m_trainingBatch;
}

size_t DeepStack::loadTrainingBatch(bool doNormalize)
{
	std::string filename = m_dir + "/" + m_name + ".training.dat";
	std::string path = m_dir + "/" + m_name + ".training.dat";
	bool success = m_trainingBatch.load(filename, arma::arma_ascii);

	if (success)
	{
		if (doNormalize)
		{
			m_trainingBatch = Rbm::normalize(m_trainingBatch);
		}
		std::cout << "Loaded " << m_trainingBatch.n_rows << " training samples\n";
	}
	
	// Migrate context part to training data
	size_t numTraining = m_trainingBatch.n_rows;
	size_t numVisible = getLayer(0)->numVisible();
	
	if (m_trainingBatch.n_cols < numVisible)
	{
		size_t diff = numVisible - m_trainingBatch.n_cols;
		arma::mat training_with_ctx = arma::join_rows(m_trainingBatch, arma::zeros(numTraining, diff));
		m_trainingBatch = training_with_ctx;
	}
	else if (m_trainingBatch.n_cols > numVisible)
	{
		m_trainingBatch = m_trainingBatch.submat(0, 0, numTraining-1, numVisible-1);
	}

	return m_trainingBatch.n_rows;
}

size_t DeepStack::saveTrainingBatch()
{
	std::string filename = m_dir + "/" + m_name + ".training.dat";
	bool success = m_trainingBatch.save(filename, arma::arma_ascii);

	if (success)
	{
		std::cout << "Saved " << m_trainingBatch.n_rows << " training samples\n";
	}
	
	return m_trainingBatch.n_rows;
}

size_t DeepStack::numTraining()
{
	return m_trainingBatch.n_rows;
}

void DeepStack::addTraining(const arma::mat &toAdd)
{ 
	m_trainingBatch.insert_rows(m_trainingBatch.n_rows, toAdd);
}

void DeepStack::delTraining(int index)
{
	m_trainingBatch.shed_row(index);
}

arma::mat DeepStack::upPass(size_t layerId, const arma::mat& v)
{
	arma::mat h = arma::zeros(0,0);
	arma::mat tv = v;
	Layer *pLayer = getLayer(layerId);
	while(pLayer)
	{
		if (pLayer->isEnable())
		{
			h = pLayer->to_h_gibbs(tv);
			tv = h;
		}
		pLayer = pLayer->next;
	}
	return h;
}

arma::mat DeepStack::downPass(size_t layerId, const arma::mat& h)
{
	arma::mat v = arma::zeros(0,0);
	arma::mat th = h;
	Layer *pLayer = getLayer(layerId);
	while(pLayer)
	{
		if (pLayer->isEnable())
		{
			v = pLayer->to_v_gibbs(th);
			th = v;
		}
		pLayer = pLayer->prev;
	}
	return v;
}

arma::mat DeepStack::upDownPass(size_t layerId, const arma::mat& v)
{
	arma::mat h = upPass(layerId, v);
	arma::mat r = downPass(numLayers()-1, h);
	return r;
}