/*
 * 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:   Stack.cpp
 * Author: jens
 * 
 * Created on 25. Oktober 2019, 18:26
 */

#include "Stack.hpp"

using namespace std;

Stack::Stack(const std::string &prjname)
: m_prjname(prjname)
, m_pLayers(nullptr)
{
}

Stack::Stack(const Stack& orig)
: m_prjname(orig.m_prjname)
, m_pLayers(orig.m_pLayers)
{
}

Stack::~Stack()
{
}

void Stack::addLayer(RbmLayer *pOtherLayer)
{
	if (!m_pLayers)
	{
		m_pLayers = pOtherLayer;
		pOtherLayer->lower = nullptr;
	}
	else
	{
		RbmLayer *pLayer = m_pLayers;
		while(pLayer->upper) 
		{
			pLayer = pLayer->upper;
		}
		pLayer->upper = pOtherLayer;
		pOtherLayer->lower = pLayer;
	}
}

RbmLayer* Stack::getLayer(size_t layerId) const
{
	RbmLayer *pLayer = m_pLayers;
	while(pLayer)
	{
		if (pLayer->id() == layerId)
		{
			return pLayer;
		}
		pLayer = pLayer->upper;
	}
	return nullptr;
	
}

void Stack::save(size_t numTraining)
{
	std::cout << "Exporting Project " << m_prjname << std::endl;
	ofstream ofs(m_prjname + string(".prj"));
	
	Json::StyledWriter writer;
	Json::Value project;
	project["stack"]["name"] = m_prjname;
	project["stack"]["num_training"] = (int)numTraining;
	project["stack"]["training_file"] = m_prjname + string(".training.dat");
	
	Json::Value layers(Json::arrayValue);
	RbmLayer *pLayer = m_pLayers;
	while(pLayer)
	{
		layers.append(pLayer->toJson());
		pLayer = pLayer->upper;
	}
	project["stack"]["layers"] = layers;
	
	ofs << writer.write(project);
}

void Stack::saveWeights()
{
	RbmLayer *pLayer = m_pLayers;
	while(pLayer)
	{
		pLayer->saveWeights();
		pLayer = pLayer->upper;
	}
}

void Stack::train(const arma::mat& batch, size_t miniBatchSize, size_t numEpochs, Rbm::IListener* pListener)
{
	RbmLayer *pLayer = m_pLayers;
	while(pLayer)
	{
		train(pLayer->id(), batch, miniBatchSize, numEpochs, pListener);
		pLayer = pLayer->upper;
	}
}

void Stack::train(size_t layerId, const arma::mat& batch, size_t miniBatchSize, size_t numEpochs, Rbm::IListener* pListener)
{
	arma::mat thisBatch = batch;
	RbmLayer *pLayer = m_pLayers;
	while(pLayer)
	{
		if (pLayer->id() == layerId)
		{
			break;
		}
		thisBatch = pLayer->toHiddenProbs(thisBatch);
		pLayer = pLayer->upper;
	}
	
	if (pLayer)
	{
		std::cout << m_prjname << ": " << " Training of layer " << std::to_string(layerId) << std::endl;
		pLayer->train(thisBatch, miniBatchSize, numEpochs, pListener);
	}
}