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

#include "Rbm.hpp"

void mylog(const char* format, ...);
#define printf mylog

#define EPSILON_SIGMA	0.001

Rbm::Rbm(Weights &weights, const MatrixXd &batch)
: m_w(weights)
, m_batch(batch)
, m_progress(0)
{
	Noise_Init(&m_noise, 0x32727155);
	updateHiddenBatch();
}

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

void Rbm::noiseGaussian(MatrixXd &dst)
{
	for (size_t i=0; i < dst.rows(); i++)
	{
		for (size_t j=0; j < dst.cols(); j++)
		{
			dst(i, j) = Noise_Gaussian(&m_noise);
		}
	}
}

void Rbm::noiseUniform(MatrixXd &dst)
{
	for (size_t i=0; i < dst.rows(); i++)
	{
		for (size_t j=0; j < dst.cols(); j++)
		{
			dst(i, j) = Noise_Uniform(&m_noise);
		}
	}
}

void Rbm::sampleGaussian(MatrixXd &dst, MatrixXd const &src)
{
	MatrixXd n(src.rows(), src.cols());

	noiseGaussian(n);
	dst = n.array() + src.array();
}

void Rbm::sampleGaussian(MatrixXd &srcDst)
{
	sampleGaussian(srcDst, srcDst);
}

MatrixXd Rbm::sample(MatrixXd const &src)
{
	MatrixXd n(src.rows(), src.cols());

	noiseUniform(n);

	return (src.array() >= n.array()).cast<double>();
}

void Rbm::sample(MatrixXd &dst, MatrixXd const &src)
{
	dst = sample(src);
}

MatrixXd Rbm::probsLogistic(MatrixXd const &src)
{
	return (1 + (-src.array()).exp()).array().cwiseInverse();
}

RowVectorXd Rbm::probsLogistic(RowVectorXd const &src)
{
	return (1 + (-src.array()).exp()).array().cwiseInverse();
}

MatrixXd Rbm::normalizeData(MatrixXd const &src)
{
#if 1
	double mean = src.array().mean();
	cout << "mean" << ": " << endl << mean << endl;

	MatrixXd x = src - mean*MatrixXd::Ones(src.rows(), src.cols());
	MatrixXd x2 = x.cwiseProduct(x);
	double stddev = sqrt(x2.array().mean());
	cout << "stddev" << ": " << endl << stddev << endl;
	
	return x/stddev;
	
#else
	MatrixXd mean = src.colwise().mean();
	MatrixXd x = src - mean.replicate(src.rows(), 1);
	MatrixXd x2 = x.cwiseProduct(x);
	
	
	MatrixXd stddev_inv = x2.colwise().mean();
	stddev_inv = stddev_inv.array().sqrt().cwiseInverse();
	return x*stddev_inv.replicate(src.rows(), 1);
	
#endif
	
}

void Rbm::train(size_t numEpochs, size_t miniBatchSize, bool &doStop)
{
	size_t i;
	size_t epoch;
	size_t gibbs;

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

	MatrixXd grad_bias_v(MatrixXd::Zero(1, m_w.getNumVisible()));
	MatrixXd grad_bias_h(MatrixXd::Zero(1, m_w.getNumHidden()));
	MatrixXd grad_weight(MatrixXd::Zero(m_w.getNumVisible(), m_w.getNumHidden()));
	MatrixXd __batch = m_batch;
	MatrixXd momentum_weights = MatrixXd::Zero(m_w.getNumVisible(), m_w.getNumHidden());
	MatrixXd momentum_bias_v(MatrixXd::Zero(1, m_w.getNumVisible()));
	MatrixXd momentum_bias_h(MatrixXd::Zero(1, m_w.getNumHidden()));
	MatrixXd penalty_weights = MatrixXd::Zero(m_w.getNumVisible(), m_w.getNumHidden());
	double L1 = 0;
	double L2 = 0;
	
	m_progress = 0;
	while (trainingSizeRemain)
	{
		if (doStop)
		{
			break;
		}
		
		cout << "trainingSizeRemain: " << trainingSizeRemain << endl;
		size_t toSlice = std::min(miniBatchSize, trainingSizeRemain);
		MatrixXd batch = __batch.block(batchRowIndex, 0, toSlice, m_w.getNumVisible());
		trainingSizeRemain -= toSlice;
		batchRowIndex += toSlice;
		size_t batchSize = batch.rows();
		double learning_rate = m_params.m_learningRate/std::min(miniBatchSize, trainingSize);
		double weight_decay = m_params.m_weightDecay/std::min(miniBatchSize, trainingSize);

		MatrixXd vis_state(batchSize, m_w.getNumVisible());
		MatrixXd vis_probs(batchSize, m_w.getNumVisible());
		MatrixXd hid_state(batchSize, m_w.getNumHidden());
		MatrixXd hid_probs(batchSize, m_w.getNumHidden());

		for (epoch=0; epoch < numEpochs; epoch++)
		{
			if (doStop)
			{
				break;
			}
			onProgressChanged();

			// Create hidden layer base on training data
			if (m_params.m_doSampleBatch)
			{
				// When the hidden units are being driven by data, always use stochastic binary states
				vis_state = sample(batch);
			}
			else
			{
				vis_state = batch;
			}
			
			hid_state = vis_state * m_w.weights() + m_w.hiddenBias().replicate(batchSize, 1);
			hid_probs = probsLogistic(hid_state);

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

			// Update weights (positive phase)
			grad_weight = vis_state.transpose() * hid_state;
			grad_bias_v = vis_state.colwise().sum();
			grad_bias_h = hid_state.colwise().sum();

			for (gibbs=0; gibbs < m_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.weights().transpose() + m_w.visibleBias().replicate(batchSize, 1);
				vis_probs = probsLogistic(vis_state);
				if (m_params.m_doSampleVisible)
				{
					vis_state = sample(vis_probs);
				}
				else
				{
					vis_state = vis_probs;
				}

				// Create hidden representation given v
				hid_state = vis_state * m_w.weights() + m_w.hiddenBias().replicate(batchSize, 1);
				hid_probs = probsLogistic(hid_state);
			
			}

			// Update weights (negative phase)
			grad_bias_v -= vis_probs.colwise().sum();
			grad_bias_h -= hid_probs.colwise().sum();
			grad_weight -= vis_probs.transpose() * hid_probs;

			for (int i=0; i < m_w.weights().rows(); i++)
			{
				for (int j=0; j < m_w.weights().cols(); j++)
				{
					if (m_w.weights()(i,j) >= 0)
					{
						penalty_weights(i,j) = weight_decay;
					}
					else
					{
						penalty_weights(i,j) = -weight_decay;
					}
				}
			}
			
			L1 = m_w.weights().array().abs().sum();
			L2 = m_w.weights().array().square().sum();
			momentum_bias_v = m_params.m_momentum*momentum_bias_v + grad_bias_v;
			momentum_bias_h = m_params.m_momentum*momentum_bias_h + grad_bias_h;
			momentum_weights = m_params.m_momentum*momentum_weights + grad_weight - L2*penalty_weights;
			
			m_w.visibleBias() += learning_rate*momentum_bias_v;
			m_w.hiddenBias() += learning_rate*momentum_bias_h;
			m_w.weights() += learning_rate*momentum_weights;
			
			m_progress += dProgress;

		} // Number of epochs
			
		MatrixXd diffErr = batch - vis_probs;
		diffErr.array() *= diffErr.array();
		double err = diffErr.colwise().sum().mean();
		cout << "error (per mini batch) = " << err << endl;
		cout << "L1                     = " << L1 << endl;
		cout << "L2                     = " << L2 << endl;
	} // number of mini batches
	
	updateHiddenBatch();
	
	MatrixXd vis = m_h * m_w.weights().transpose();
	vis += m_w.visibleBias().replicate(__batch.rows(), 1);
	MatrixXd diffErr = __batch - probsLogistic(vis);
	diffErr.array() *= diffErr.array();
	double err = diffErr.colwise().sum().mean();
	cout << "error (total) = " << err << endl;

	onProgressChanged();
	
}

double Rbm::getProgress() const
{
	return m_progress;
}

void Rbm::toHidden(RowVectorXd &h, RowVectorXd const &v)
{
	h = v * m_w.weights();
	h += m_w.hiddenBias();
	h = probsLogistic(h);
}

void Rbm::toVisible(RowVectorXd &v, RowVectorXd const &h)
{
	v = h * m_w.weights().transpose();
	v += m_w.visibleBias();
	v = probsLogistic(v);
}

void Rbm::setWeightDecay(double value)
{
	m_params.m_weightDecay = value;
	onParamsChanged();
}

void Rbm::setDoRaoBlackwell(bool flag)
{
	m_params.m_doRaoBlackwell = flag;
	onParamsChanged();
}

void Rbm::setDoSampleVisible(bool flag)
{
	m_params.m_doSampleVisible = flag;
	onParamsChanged();
}

void Rbm::setDoSampleBatch(bool flag)
{
	m_params.m_doSampleBatch = flag;
	onParamsChanged();
}

void Rbm::setNumGibbs(size_t value)
{
	m_params.m_numGibbs = value;
	onParamsChanged();
}

void Rbm::setMuWeights(double value)
{
	m_params.m_learningRate = value;
	onParamsChanged();
}

void Rbm::setMomentum(double value)
{
	m_params.m_momentum = value;
	onParamsChanged();
}

MatrixXd const& Rbm::getHiddenBatch()
{
	return m_h;
}

MatrixXd const& Rbm::getBatch()
{
	return m_batch;
}

void Rbm::updateHiddenBatch()
{
	if (m_batch.rows() == 0)
	{
		return;
	}
	
	m_h.resize(m_batch.rows(), m_w.getNumHidden());
	m_h = m_batch * m_w.weights();
	m_h += m_w.hiddenBias().replicate(m_batch.rows(), 1);
	m_h = probsLogistic(m_h);
}

Rbm::Params const& Rbm::params()
{
	return m_params;
}