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

#include "Layer.hpp"
using namespace std;

Layer::Layer(const string &name, size_t id, size_t numVisibleX, size_t numVisibleY, size_t numHidden, size_t numContext)
: Rbm(numVisibleX*numVisibleY+numContext, numHidden)
, next(nullptr)	
, prev(nullptr)	
, m_name(name)
, m_id(id)
, m_numVisibleX(numVisibleX)
, m_numVisibleY(numVisibleY)
, m_numContext(numContext)
, m_context(0, numContext)
{
	cout << "Create Layer " << m_name << "." << to_string((int)m_id) << endl;
}

Layer::Layer(const Layer& orig)
: Rbm(orig.bv().n_elem, orig.bh().n_elem)
, next(nullptr)	
, prev(nullptr)	
, m_name(orig.m_name)
, m_id(orig.m_id)
, m_numVisibleX(orig.m_numVisibleX)
, m_numVisibleY(orig.m_numVisibleY)
{
}

Layer::~Layer()
{
}

size_t Layer::id()
{
	return m_id;
}

std::string& Layer::name()
{
	return m_name;
}

int Layer::numVisibleX()
{
	return m_numVisibleX;
}

int Layer::numVisibleY()
{
	return m_numVisibleY;
}

const arma::mat& Layer::context() const
{
    return m_context;
}

Layer* Layer::root()
{
	Layer *pLayer = this;
	while (pLayer->prev)
	{
		pLayer = pLayer->prev;
	}
	return pLayer;
}

arma::mat Layer::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 Layer::downPass(const arma::mat& h)
{
	arma::mat v = prob(h_to_v(h));
	if (prev)
	{
		return prev->downPass(v);
	}
	return v;
}

arma::mat Layer::upPass(const arma::mat& v)
{
	arma::mat r = v;
	arma::mat h = gibbsPass(r);
	if (next)
	{
		return next->upPass(h);
	}
	return h;
}

arma::mat Layer::upDownPass(const arma::mat& v)
{
	arma::mat r = v;
	arma::mat h = gibbsPass(r);
	if (next)
	{
		next->upDownPass(h);
	}
	else if (prev)
	{
		prev->downPass(r);
	}
	return r;
}

arma::mat Layer::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 Layer::vc_to_v(const arma::mat& vc) const
{
	return arma::reshape(vc, 1, numVisible() - m_numContext);
}

void Layer::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 v = batch.row(i - 1);
			arma::mat h = arma::zeros(1, numHidden());
			gibbs_vh(v, h);
			arma::mat training_with_ctx = arma::join_rows(batch.row(i).cols(0, numVisible() - m_numContext - 1), h);
			batch.row(i) = training_with_ctx;
		}
	}
}

void Layer::train(const arma::mat& batch, IListener* pListener)
{
	Rbm::train(trainingData(batch), pListener);
}

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

bool Layer::weightsLoad(std::string const &dir, std::string const &prj)
{
	arma::mat w;
	arma::mat bh;
	arma::mat bv;
	bool result = true;
	result &= w.load(filePrefix(dir, prj) + ".w.dat", arma::arma_ascii);
	result &= bh.load(filePrefix(dir, prj) + ".bh.dat", arma::arma_ascii);
	result &= bv.load(filePrefix(dir, prj) + ".bv.dat", arma::arma_ascii);

	if (result)
	{
		std::cout << "Layer " << m_id << ": Importing weights" << std::endl;
		weightsAssign(w, bh, bv);
	}

	return result;
}

bool Layer::weightsSave(std::string const &dir, std::string const &prj)
{
	bool result = true;
	result &= whv().save(filePrefix(dir, prj) + ".w.dat", arma::arma_ascii);
	result &= bh().save(filePrefix(dir, prj) + ".bh.dat", arma::arma_ascii);
	result &= bv().save(filePrefix(dir, prj) + ".bv.dat", arma::arma_ascii);

	if (result)
	{
		std::cout << "Layer " << m_id << ": Exporting weights" << std::endl;
	}

	return result;
}

Json::Value Layer::toJson() const
{
	std::cout << "Exporting Layer " << to_string((int)m_id) << std::endl;
	Json::Value layer;
	layer["name"] = m_name;
	layer["id"] = (int)m_id;
	layer["numVisibleX"] = (int)m_numVisibleX;
	layer["numVisibleY"] = (int)m_numVisibleY;
	layer["numHidden"] = (int)whv().n_cols;
	layer["numContext"] = m_numContext;
	layer["rbm"] = Rbm::toJson();
	
	return layer;

}

std::string Layer::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;
}