/*
 * 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:   Layer.hpp
 * Author: jens
 *
 * Created on 25. Oktober 2019, 08:13
 */

#ifndef RBMLAYER_HPP
#define RBMLAYER_HPP

#include <string>

#include <armadillo>
#include <jsoncpp/json/json.h>
#include "Rbm.hpp"

class Layer : public Rbm
{
public:
	Layer *next;
	Layer *prev;

	Layer(const std::string &name, size_t id, size_t numVisibleX, size_t numVisibleY, size_t numHidden, size_t numContext=0);
	Layer(const Layer& orig);
	virtual ~Layer();

	Json::Value toJson() const;
	bool loadWeights(const std::string &prjname="");
	bool saveWeights(const std::string &prjname="");
	
	void setBatch(arma::mat const &batch)
	{
		if (batch.n_rows > 0)
		{	
			if (m_numContext)
			{
				arma::mat c_states = arma::zeros(batch.n_rows, m_numContext);
				arma::mat ctx = arma::zeros(1, m_numContext);
				for (int i=1; i < batch.n_rows; i++)
				{
					arma::mat h = prob(v_to_h(arma::join_rows(batch.row(i-1), ctx)));
					ctx = h;
					c_states.row(i) = ctx;
				}
				arma::mat batch_with_ctx = arma::join_rows(batch, c_states);
				Rbm::setBatch(trainingData(batch_with_ctx));
			}
			else
			{
				Rbm::setBatch(trainingData(batch));
			}
		}
	}
	
	std::string& name()
	{
		return m_name;
	}
	
	size_t id()
	{
		return m_id;
	}
	
	int numVisibleX()
	{
		return m_numVisibleX;
	}

	int numVisibleY()
	{
		return m_numVisibleY;
	}
	
	int numHidden()
	{
		return whv().n_cols;
	}

	arma::mat trainingData(arma::mat const &batch)
	{
		arma::mat thisBatch = batch;
		Layer *pLayer = root();
		while (pLayer)
		{
			if (pLayer == this)
			{
				break;
			}
			thisBatch = pLayer->toHiddenProbs(thisBatch);
			pLayer = pLayer->next;
		}
		return thisBatch;
	}
		
	Layer *root()
	{
		Layer *pLayer = this;
		while(pLayer->prev)
		{
			pLayer = pLayer->prev;
		}
		return pLayer;
	}

	arma::mat vc_to_v(const arma::mat &vc) const
	{
		return arma::reshape(vc, 1, numVisible() - m_numContext);
	}

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

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

private:
	
	std::string m_name;
	std::string m_weightsFile;
	size_t m_id;
	size_t m_numVisibleX;
	size_t m_numVisibleY;
	size_t m_numContext;
	arma::mat m_context;

};

#endif /* RBMLAYER_HPP */