/*
 * 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"
#include "RnnTextHelper.hpp"

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

RnnStack::~RnnStack()
{
}

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

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_v(const arma::mat& vc, size_t numContext) const
{
	size_t numVisible = vc.n_cols;
	return vc.submat(0, 0, vc.n_rows - 1, numVisible - numContext - 1);
}

void RnnStack::train(const arma::mat& batch, Rbm::IListener* pListener)
{
	// thisBatch = {padding | batch}
	int numTraining = batch.n_rows;
	arma::mat padding = arma::zeros(getSeqLen()-1, batch.n_cols);
	arma::mat batch_padded = arma::join_cols(batch, padding);
	arma::mat c = arma::zeros(numTraining, getLayer(0)->numContext());
	
	for (int i=0; i < getSeqLen(); i++)
	{
		Layer *pLayer = getLayer(i);
		int k = i;
		arma::mat batch_shifted = batch_padded.rows(k, numTraining+k-1);
		arma::mat vc = v_to_vc(batch_shifted, c);
		pLayer->train(vc, pListener);
		pLayer->gibbs_vh(vc, c);
		pLayer = pLayer->next;
	}
}

arma::mat RnnStack::step_forward(arma::mat &state, const arma::mat& v_curr)
{
	arma::mat c = arma::zeros(1, getLayer(0)->numContext());
	arma::mat z = arma::zeros(arma::size(v_curr));
	arma::mat r;
	
	if (state.is_empty())
	{
		state = arma::zeros(getSeqLen(), v_curr.n_cols);
	}
	state = arma::shift(state, 1, 0);
	state.row(0) = v_curr;
	
	for (int j=0; j < getSeqLen(); j++)
	{
		int k = getSeqLen()- j - 1;
		Layer *pLayer = getLayer(j);
		arma::mat v = arma::join_rows(z, state.row(k));
		arma::mat vc = v_to_vc(v, c);
		pLayer->gibbs_vh(vc, c);
		
		r = vc_to_v(vc, pLayer->numContext());
	}
	return r;
}