/*
 * 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 "Rbm.hpp"
#include "noise.h"

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

Rbm::Rbm(const Rbm& orig)
{
}

Rbm::~Rbm()
{
}

void Rbm::train(const arma::mat& batch, size_t numEpochs, size_t miniBatchSize, const Params& params, IListener* pListener)
{
	size_t epoch;
	size_t gibbs;

	size_t trainingSize = batch.n_rows;
	size_t trainingSizeRemain = trainingSize;
	size_t batchRowIndex = 0;
	
	double dProgress = 1.0/(numEpochs*(double)trainingSize/std::min(miniBatchSize, trainingSize));

	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);
	
	Status status;
	
	status.progress = 0;
	while (trainingSizeRemain)
	{
		status.trainingSizeRemain = trainingSizeRemain;
		size_t toSlice = std::min(miniBatchSize, trainingSizeRemain);
		arma::mat miniBatch = batch.rows(batchRowIndex, batchRowIndex+toSlice-1);
		trainingSizeRemain -= toSlice;
		batchRowIndex += toSlice;
		size_t miniBatchSizeActual = miniBatch.n_rows;
		double learning_rate = params.m_learningRate/std::min(miniBatchSizeActual, trainingSize);
		double weight_decay = params.m_weightDecay/std::min(miniBatchSizeActual, trainingSize);

		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 < numEpochs; epoch++)
		{

			// Create hidden layer base on training data
			if (params.m_doSampleBatch)
			{
				// When the hidden units are being driven by data, always use stochastic binary states
				vis_state = sample(miniBatch);
			}
			else
			{
				vis_state = miniBatch;
			}
			
			hid_state = vis_state * m_w + arma::repmat(m_bh, miniBatchSizeActual, 1);
			hid_probs = probsLogistic(hid_state);

			// Sample hidden
			if (params.m_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 < params.m_numGibbs; gibbs++)
			{

				// Create hidden representation given v
				hid_state = sample(hid_probs);

				// Create visible reconstruction (a fantasy...) given hid
				vis_state = hid_state * m_w.t() + arma::repmat(m_bv, miniBatchSizeActual, 1);
				vis_probs = probsLogistic(vis_state);
				if (params.m_doSampleVisible)
				{
					vis_state = sample(vis_probs);
				}
				else
				{
					vis_state = vis_probs;
				}

				// Create hidden representation given v
				hid_state = vis_state * m_w + repmat(m_bh, miniBatchSizeActual, 1);
				hid_probs = probsLogistic(hid_state);
			
			}

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

			penalty_weights = weight_decay*arma::sign(m_w);
			
			status.L1 = accu(abs(m_w));
			status.L2 = accu(m_w % m_w);
			momentum_bias_v = params.m_momentum*momentum_bias_v + grad_bias_v;
			momentum_bias_h = params.m_momentum*momentum_bias_h + grad_bias_h;
			momentum_weights = params.m_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;
			
			status.progress += dProgress;
			
		} // Number of epochs
			
		status.epoch = epoch;
		arma::mat diffErr = miniBatch - vis_probs;
		arma::mat diffErr_squared = diffErr % diffErr;
		status.err = accu(diffErr_squared);

		if (pListener)
		{
			if(!pListener->onProgress(status))
			{
				break;
			}
		}
		
	} // number of mini batches
}

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;
	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) >= Noise_Uniform(&m_noise);
		}
	}
	return dst;
}

arma::mat Rbm::toHidden(const arma::mat &visible)
{
	arma::mat h = visible.t() * m_w  + m_bh;
	return h;
}

arma::mat Rbm::toVisible(const arma::mat &hidden)
{
	arma::mat v = hidden * m_w.t() + m_bv;
	return v;
}

void Rbm::uniform(arma::mat& srcDst, double mu, double stdDev)
{
	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;
		}
	}
}