/*
 * 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:   RnnStack.cpp
 * Author: jens
 * 
 * Created on 17. Januar 2022, 12:43
 */

#include "RnnStack.hpp"

RnnStack::RnnStack(const std::string &name, size_t numContext)
: AStack(StackType::Rnn, name)
, m_numContext(numContext)
{
}

RnnStack::~RnnStack()
{
}

size_t RnnStack::getSeqLen()
{
	return numLayers();
}

size_t RnnStack::numContext()
{
	return m_numContext;
}

arma::mat RnnStack::v_to_vc(const arma::mat& v) const
{
	arma::mat c = arma::zeros(v.n_rows, m_numContext);
	return arma::join_rows(v, c);
}

arma::mat RnnStack::v_to_vc(const arma::mat& v, const arma::mat& c) const
{
	return arma::join_rows(v, c);
}

arma::mat RnnStack::vc_to_c(const arma::mat& vc) const
{
	size_t numVisible = vc.n_cols;
	
	if (m_numContext == 0)
	{
		return arma::mat(1, 0);
	}
	return vc.submat(0, numVisible - m_numContext, 0, numVisible - 1);
}

arma::mat RnnStack::vc_to_v(const arma::mat& vc) const
{
	size_t numVisible = vc.n_cols;
	return vc.submat(0, 0, vc.n_rows - 1, numVisible - m_numContext - 1);
}

arma::mat RnnStack::trainingBatchFrom(size_t layerId, const arma::mat& batch)
{
	arma::mat vc = batch;
	Layer *pLayer = getLayer(0);
	while (pLayer)
	{
		if (layerId == pLayer->id())
		{
			break;
		}
		arma::mat c = pLayer->to_h_gibbs(vc);
		arma::mat v = arma::shift(vc_to_v(vc), layerId, 1);
		vc = v_to_vc(v, c);
		pLayer = pLayer->next;
	}
	return vc;
}

void RnnStack::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 RnnStack::train(size_t layerId, const arma::mat& batch, Rbm::IListener* pListener)
{
	arma::mat thisBatch = batch;
	Layer *pLayer = getLayer(layerId);
	if (pLayer)
	{
		thisBatch = trainingBatchFrom(pLayer->id(), batch);
		pLayer->train(thisBatch, pListener);
	}
}