/*
 * 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:   Rbm.cpp
 * Author: jens
 * 
 * Created on 21. Oktober 2019, 21:28
 */

#include <cassert>
#include "Rbm.hpp"

Rbm::Rbm(size_t numVisible, size_t numHidden)
: m_params()
, m_w(numVisible, numHidden)
, m_bv(1, numVisible)
, m_bh(1, numHidden)
{
	assert(numVisible > 0);
	assert(numHidden > 0);
	Noise_Init(&m_noise, 0x32727155);
}

Rbm::Rbm(const Rbm& orig)
: m_params(orig.m_params)
, m_w(orig.m_w)
, m_bv(orig.m_bv)
, m_bh(orig.m_bh)
{
}

Rbm::~Rbm()
{
	Noise_Free(&m_noise);
}

void Rbm::weightsInit(double stddev, double mu)
{
	uniform(m_w, stddev, mu);
	uniform(m_bv, stddev, mu);
	uniform(m_bh, stddev, mu);

}

void Rbm::fromJson(Json::Value rbm)
{
	std::cout << "Importing Rbm" << std::endl;
	m_params.fromJson(rbm["params"]);
}

Json::Value Rbm::toJson() const
{
	std::cout << "Exporting Rbm" << std::endl;
	Json::Value rbm;
	rbm["params"] = m_params.toJson();
	return rbm;
}

void Rbm::train(const arma::mat& batch, IListener* pListener)
{
	Status status;
	size_t epoch;
	size_t gibbs;

	status.trainingSizeRemain = batch.n_rows;
	size_t batchRowIndex = 0;
	
	double dProgress = 1.0/status.trainingSizeRemain;

	arma::mat grad_bias_v(arma::zeros(1, m_w.n_rows));
	arma::mat grad_bias_h(arma::zeros(1, m_w.n_cols));
	arma::mat grad_weight(arma::zeros(m_w.n_rows, m_w.n_cols));
	arma::mat momentum_weights = arma::zeros(m_w.n_rows, m_w.n_cols);
	arma::mat momentum_bias_v(arma::zeros(1, m_w.n_rows));
	arma::mat momentum_bias_h(arma::zeros(1, m_w.n_cols));
	arma::mat penalty_weights = arma::zeros(m_w.n_rows, m_w.n_cols);
	
	if (pListener)
	{
		pListener->onProgress(this, status);
	}

	while (status.trainingSizeRemain)
	{
		size_t miniBatchSizeActual = std::min(m_params.miniBatchSize, status.trainingSizeRemain);
		arma::mat miniBatch = batch.rows(batchRowIndex, batchRowIndex+miniBatchSizeActual-1);
		status.trainingSizeRemain -= miniBatchSizeActual;
		batchRowIndex += miniBatchSizeActual;
		double learning_rate = m_params.learningRate/miniBatchSizeActual;
		double weight_decay = m_params.weightDecay/miniBatchSizeActual;

		arma::mat vis_state(miniBatchSizeActual, m_w.n_rows);
		arma::mat vis_probs(miniBatchSizeActual, m_w.n_rows);
		arma::mat hid_state(miniBatchSizeActual, m_w.n_cols);
		arma::mat hid_probs(miniBatchSizeActual, m_w.n_cols);

		for (epoch=0; epoch < m_params.numEpochs; epoch++)
		{

			// Create hidden layer base on training data
			if (m_params.doSampleBatch)
			{
				// When the hidden units are being driven by data, always use stochastic binary states
				vis_state = sample(miniBatch);
			}
			else
			{
				vis_state = miniBatch;
			}
			
			hid_probs = probsLogistic(toHiddenState(vis_state));

			// Sample hidden
			if (m_params.doRaoBlackwell)
			{
				hid_state = hid_probs;
			}
			else
			{
				hid_state = sample(hid_probs);
			}

			// Update weights (positive phase)
			grad_weight = vis_state.t() * hid_state;
			grad_bias_v = sum(vis_state, 0);
			grad_bias_h = sum(hid_state, 0);

			for (gibbs=0; gibbs < m_params.numGibbs; gibbs++)
			{
				// Create visible reconstruction (a fantasy...) given hid
				if (m_params.gibbsDoSampleHidden)
				{
					vis_probs = toVisibleProbs(sample(hid_probs));
				}
				else
				{
					vis_probs = toVisibleProbs(hid_probs);
				}

				// Create hidden representation given v
				if (m_params.gibbsDoSampleVisible)
				{
					hid_state = toHiddenState(sample(vis_probs));
				}
				else
				{
					hid_state = toHiddenState(vis_probs);
				}
				hid_probs = probsLogistic(hid_state);
			
			}

			// Update weights (negative phase)
			grad_weight -= vis_probs.t() * hid_probs;
			grad_bias_v -= sum(vis_probs, 0);
			grad_bias_h -= sum(hid_probs, 0);

			penalty_weights = weight_decay*arma::sign(m_w);
			
			status.L1 = accu(abs(m_w));
			status.L2 = accu(m_w % m_w);
			momentum_bias_v = m_params.momentum*momentum_bias_v + grad_bias_v;
			momentum_bias_h = m_params.momentum*momentum_bias_h + grad_bias_h;
			momentum_weights = m_params.momentum*momentum_weights + grad_weight - status.L2*penalty_weights;
			
			m_bv += learning_rate*momentum_bias_v;
			m_bh += learning_rate*momentum_bias_h;
			m_w += learning_rate*momentum_weights;
			
		} // Number of epochs
			
		status.epoch = epoch;
		arma::mat diffErr = miniBatch - vis_probs;
		arma::mat diffErr_squared = diffErr % diffErr;
		status.err = accu(diffErr_squared)/diffErr_squared.n_elem;
		status.progress += dProgress*miniBatchSizeActual;
		
		if (pListener)
		{
			if(!pListener->onProgress(this, status))
			{
				break;
			}
		}
		
	} // number of mini batches

	arma::mat hid_probs = toHiddenProbs(batch);
	arma::mat vis_probs = toVisibleProbs(hid_probs);
	arma::mat diffErr = batch - vis_probs;
	arma::mat diffErr_squared = diffErr % diffErr;
	status.err_total = accu(diffErr_squared)/diffErr_squared.n_elem;

	if (pListener)
	{
		pListener->onProgress(this, status);
	}
}

arma::mat Rbm::probsLogistic(const arma::mat &src)
{
	return 1 / (1 + (arma::exp(-src)));
}

arma::mat Rbm::sample(const arma::mat &src)
{
	arma::mat dst = src;
	uniform(dst);
	
	for (size_t i=0; i < src.n_rows; i++)
	{
		for (size_t j=0; j < src.n_cols; j++)
		{
			dst(i, j) = src(i, j) >= dst(i, j);
		}
	}
	return dst;
}

arma::mat Rbm::toHiddenState(const arma::mat &visible) const
{
	return visible * m_w  + arma::repmat(m_bh, visible.n_rows, 1);
}

arma::mat Rbm::toVisibleState(const arma::mat &hidden) const
{
	return hidden * m_w.t() + arma::repmat(m_bv, hidden.n_rows, 1);
}

arma::mat Rbm::toHiddenProbs(const arma::mat &visible) const
{
	return probsLogistic(toHiddenState(visible));
}

arma::mat Rbm::toVisibleProbs(const arma::mat &hidden) const
{
	return probsLogistic(toVisibleState(hidden));
}

void Rbm::uniform(arma::mat& srcDst, double stdDev, double mu)
{
#if 1
	for (size_t i=0; i < srcDst.n_rows; i++)
	{
		for (size_t j=0; j < srcDst.n_cols; j++)
		{
			srcDst(i, j) = stdDev*(Noise_Uniform(&m_noise) + mu - 0.5);
		}
	}
#else
	srcDst = stdDev*(arma::randu(srcDst.n_rows, srcDst.n_cols) + mu - 0.5);
#endif
}

const arma::mat& Rbm::w() const
{
	return m_w;
}

const arma::mat& Rbm::bv() const
{
	return m_bv;
}

const arma::mat& Rbm::bh() const
{
	return m_bh;
}