/*
 * 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;

	void calcContextBatch(arma::mat &batch)
	{
		size_t numTraining = batch.n_rows;

		if (m_numContext > 0 and numTraining > 0)
		{
			for (int i=1; i < numTraining; i++)
			{
				arma::mat h = toHiddenProbs(batch.row(i-1));
				arma::mat training_with_ctx = arma::join_rows(batch.row(i).cols(0, numVisible()-m_numContext-1), h);
				batch.row(i) = training_with_ctx;
			}
		}
	}
	
	std::string& name()
	{
		return m_name;
	}
	
	size_t id()
	{
		return m_id;
	}
	
	int numVisibleX()
	{
		return m_numVisibleX;
	}

	int numVisibleY()
	{
		return m_numVisibleY;
	}
	
	arma::mat gibbsPass(arma::mat &vr)
	{
		arma::mat h;
		for (int i=0; i < params().numGibbs; i++)
		{
			h = prob(v_to_h(vr));
			vr = prob(h_to_v(h));
		}
		return h;
	}
	
	arma::mat upPass(arma::mat const &v)
	{
		arma::mat r = v;
		arma::mat h = gibbsPass(r);
		if (next)
		{
			next->upPass(h);
		}
		else
		{
			return h;
		}
	}
	
	arma::mat downPass(arma::mat const &h)
	{
		arma::mat v = prob(h_to_v(h));
		if (prev)
		{
			return prev->downPass(v);
		}
		return v;
	}
	
	arma::mat upDownPass(arma::mat const &v)
	{
		arma::mat r = v;
		arma::mat h = gibbsPass(r);
		if (next)
		{
			next->upDownPass(h);
		}
		else if (prev)
		{
			prev->downPass(r);
		}
		return r;
	}

	bool weightsLoad(std::string const &dir, std::string const &prj);
	bool weightsSave(std::string const &dir, std::string const &prj);
	
	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);
	}

	const arma::mat& context() const
	{
		return m_context;
	}
	
private:
	
	std::string m_name;
	size_t m_id;
	size_t m_numVisibleX;
	size_t m_numVisibleY;
	size_t m_numContext;
	arma::mat m_context;
	
	// Compatibility
	bool loadWeights(const std::string &prjname="");
	bool saveWeights(const std::string &prjname="");
	std::string filePrefix(const std::string &dir, const std::string &prjname) const
	{
		std::string  filename = m_name + "." + std::to_string((int)m_id);
		if (prjname.size() > 0)
		{
			filename = prjname + "." + filename;
		}
		return dir + "/" + filename;
	}

};

#endif /* RBMLAYER_HPP */