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

#define RBM_TRAIN_FLAT 0

Rbm::Rbm(size_t numVisible, size_t numHidden, size_t numContext)
: m_params()
, m_whv(numVisible+numContext, numHidden)
, m_bhv(1, numHidden)
, m_bv(1, numVisible+numContext)
, m_ctx()
{
	assert(numVisible > 0);
	assert(numHidden > 0);
	if (numContext)
	{
		assert(numContext == numHidden);
		m_ctx.resize(1, numContext);
	}
	Noise_Init(&m_noise, 0x32727155);
}

Rbm::Rbm(const Rbm& orig)
: m_params(orig.m_params)
, m_whv(orig.m_whv)
, m_bhv(orig.m_bhv)
, m_bv(orig.m_bv)
, m_ctx(orig.m_ctx)
{
}

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

void Rbm::weightsInit(double stddev, double mu)
{
	uniform(m_whv, stddev, mu);
	uniform(m_bhv, stddev, mu);
	uniform(m_bv, stddev, mu);
	uniform(m_ctx, 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::gibbs(arma::mat &hv_probs, arma::mat &v_probs)
{
	for (int i=0; i < m_params.numGibbs; i++)
	{
		// Create visible reconstruction (a fantasy...) given hid
		if (m_params.gibbsDoSampleHidden)
		{
			v_probs = prob(h_to_v(sample(hv_probs)));
		}
		else
		{
			v_probs = prob(h_to_v(hv_probs));
		}

		// Create hidden representation given v
		if (m_params.gibbsDoSampleVisible)
		{
			hv_probs = prob(v_to_h(sample(v_probs)));
		}
		else
		{
			hv_probs = prob(v_to_h(v_probs));
		}
	}
}

void Rbm::weightUpdate(arma::mat const &v_states, arma::mat &dw, arma::mat &dbhv, arma::mat &dbv)
{
	arma::mat v_probs(v_states);
	arma::mat h_states = v_to_h(v_states);
	arma::mat h_probs = prob(h_states);
	
	// Sample hidden
	if (m_params.doRaoBlackwell)
	{
		h_states = h_probs;
	}
	else
	{
		h_states = sample(h_probs);
	}
	
	// Update weights (positive phase)
	dw = v_states.t() * h_states;
	dbv = sum(v_states, 0);
	dbhv = sum(h_states, 0);

	gibbs(h_probs, v_probs);

	// Update weights (negative phase)
	dw -= v_probs.t() * h_probs;
	dbv -= sum(v_probs, 0);
	dbhv -= sum(h_probs, 0);
}

void Rbm::train(const arma::mat& batch, IListener* pListener)
{
	Status status;
	double dProgress = 100.0/(batch.n_rows*m_params.numEpochs);
	double progress = 0;
	int lastProgress = -100;
	int batchRowIndex = 0;
	
	arma::mat grad_bias_v(arma::zeros(1, m_bv.n_cols));
	arma::mat grad_bias_hv(arma::zeros(1, m_bhv.n_cols));
	arma::mat grad_weight_hv(arma::zeros(m_whv.n_rows, m_whv.n_cols));
	arma::mat momentum_whv = arma::zeros(m_whv.n_rows, m_whv.n_cols);
	arma::mat momentum_bias_v(arma::zeros(1, m_bv.n_cols));
	arma::mat momentum_bias_hv(arma::zeros(1, m_bhv.n_cols));
	arma::mat penalty_weights = arma::zeros(m_whv.n_rows, m_whv.n_cols);
	arma::mat ctx = arma::zeros(1, numContext());

	int trainingSizeRemain = batch.n_rows;

	bool shouldAbort = false;
	while (trainingSizeRemain && !shouldAbort)
	{
		int miniBatchSizeActual = std::min(m_params.miniBatchSize, trainingSizeRemain);
		arma::mat miniBatch_v = batch.rows(batchRowIndex, batchRowIndex+miniBatchSizeActual-1);
		trainingSizeRemain -= miniBatchSizeActual;
		batchRowIndex += miniBatchSizeActual;
		int scaler = std::min(m_params.miniBatchSize, (int)batch.n_rows);
		double learning_rate = m_params.learningRate/scaler;
		double weight_decay = m_params.weightDecay/scaler;

#if RBM_TRAIN_FLAT
		arma::mat h_probs(miniBatchSizeActual, m_bhv.n_cols);
		arma::mat v_probs(miniBatchSizeActual, m_bv.n_cols);
		arma::mat hid_states(miniBatchSizeActual, m_bhv.n_cols);
#endif
		arma::mat miniBatch;
		if (numContext() > 0)
		{
			arma::mat c_states(arma::zeros(miniBatchSizeActual, numContext()));

			for (int i=1; i < miniBatchSizeActual; i++)
			{
				arma::mat h = prob(v_to_h(arma::join_rows(miniBatch_v.row(i-1), ctx)));
				ctx = h;
				c_states.row(i) = ctx;
			}
			miniBatch = (arma::join_rows(miniBatch_v, c_states));
		}
		else
		{
			miniBatch = miniBatch_v;
		}
		arma::mat v_states(miniBatch);

		// 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
			v_states = sample(miniBatch);
		}
		
		for (int epoch=0; epoch < m_params.numEpochs; epoch++)
		{
#if RBM_TRAIN_FLAT
			// Sample hidden
			h_probs = prob(v_to_h(v_states));
			if (m_params.doRaoBlackwell)
			{
				hid_states = h_probs;
			}
			else
			{
				hid_states = sample(h_probs);
			}
			
			// Update weights (positive phase)
			grad_weight = v_states.t() * hid_states;
			grad_bias_v = sum(v_states, 0);
			grad_bias_hv = sum(hid_states, 0);

			for (int i=0; i < m_params.numGibbs; i++)
			{
				// Create visible reconstruction (a fantasy...) given hid
				if (m_params.gibbsDoSampleHidden)
				{
					v_probs = prob(h_to_v(sample(h_probs)));
				}
				else
				{
					v_probs = prob(h_to_v(h_probs));
				}

				// Create hidden representation given v
				if (m_params.gibbsDoSampleVisible)
				{
					hid_states = v_to_h(sample(v_probs));
				}
				else
				{
					hid_states = v_to_h(v_probs);
				}
				h_probs = prob(hid_states);
			
			}

			// Update weights (negative phase)
			grad_weight -= v_probs.t() * h_probs;
			grad_bias_v -= sum(v_probs, 0);
			grad_bias_hv -= sum(h_probs, 0);
#else
			if (m_ctx.n_cols == 0)
			{
				weightUpdate(v_states, grad_weight_hv, grad_bias_hv, grad_bias_v);
			}
			else
			{
				weightUpdate(v_states, grad_weight_hv, grad_bias_hv, grad_bias_v);
			}
#endif
			
			penalty_weights = weight_decay*arma::sign(m_whv);
			
			status.L1 = accu(abs(m_whv));
			status.L2 = accu(m_whv % m_whv);
			momentum_bias_v = m_params.momentum*momentum_bias_v + grad_bias_v;
			momentum_bias_hv = m_params.momentum*momentum_bias_hv + grad_bias_hv;
			momentum_whv = m_params.momentum*momentum_whv + grad_weight_hv - status.L2*penalty_weights;
			
			m_bv += learning_rate*momentum_bias_v;
			m_bhv += learning_rate*momentum_bias_hv;
			m_whv += learning_rate*momentum_whv;
			
			progress += dProgress*miniBatchSizeActual;
			status.progress = (int)(progress + 0.5);

			if (status.progress != lastProgress)
			{
				lastProgress = status.progress;
				
				// Calculate error
				arma::mat diffErr = miniBatch - prob(h_to_v(prob(v_to_h(v_states))));
				arma::mat diffErr_squared = diffErr % diffErr;
				status.err = accu(diffErr_squared)/diffErr_squared.n_elem;
				if (pListener)
				{
					if(!pListener->onProgress(this, status))
					{
						shouldAbort = true;
						break;
					}
				}
			}
			
		} // Number of epochs
			
	} // number of mini batches

#if FIXED_BATCH
	arma::mat diffErr = batch - prob(h_to_v(prob(v_to_h(batch))));
	arma::mat diffErr_squared = diffErr % diffErr;
	status.err_total = accu(diffErr_squared)/diffErr_squared.n_elem;

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

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

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

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

arma::mat Rbm::v_to_h(const arma::mat &visible) const
{
	return visible * m_whv  + arma::repmat(m_bhv, visible.n_rows, 1);
}

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

arma::mat Rbm::normalize(const arma::mat& src)
{
	double mean = arma::accu(src)/src.n_elem;
	arma::mat x = src - mean;
	arma::mat x2 = x % x;
	double stddev = sqrt(arma::accu(x2)/x2.n_elem);
	
	std::cout << "mean" << "  : " << std::endl << mean << std::endl;
	std::cout << "stddev" << ": " << std::endl << stddev << std::endl;
	return x/stddev;
}

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::whv() const
{
	return m_whv;
}

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

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