/*
 * Rbm.hpp
 *
 *  Created on: 21.09.2014
 *      Author: jens
 */

#ifndef RBM_HPP_
#define RBM_HPP_

#include "VisibleLayer.hpp"
#include "HiddenLayer.hpp"
#include "Weights.hpp"
#include "LayerArray.hpp"
#include <cmath>
#include <Eigen/Dense>

using namespace Eigen;

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

#define EPSILON_SIGMA	0.05

class Rbm;

class RbmListener
{
public:
	RbmListener() {}
	virtual ~RbmListener()
	{
	}

	virtual void onEpochTrained(const Rbm &obj) = 0;
};

class Rbm
{
public:
	Rbm(Weights &weights, const MatrixXd &batch, RbmListener *pListener = nullptr)
	: m_w(weights)
	, m_batch(batch)
	, m_variableSigma(weights.getNumVisible())
	, m_constantSigma(1.0)
	, m_pListener(pListener)
	, m_progress(0)
	, m_sigmaDecay(1.0)
	, m_weightDecay(0.0)
	, m_lambda(1.0)
	, m_sparsity(0)
	, m_muWeights(0.01)
	, m_muSparsity(0.01)
	, m_momentum(0.5)
	, m_useVisibleGaussian(false)
	, m_doRaoBlackwell(false)
	, m_useProbsForHiddenReconstruction(false)
	, m_doSparse(false)
	, m_doNormalizeData(false)
	, m_doLearnVariance(false)
	, m_numGibbs(1)
	{
		Noise_Init(&m_noise, 0x32727155);

#if 1
		VectorXd a(4);
		a << 1, 2, 3, 4;
		VectorXd b(4);

		b.array() = -a.array().exp();

		cout << b << endl;
		
#endif
	}

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

	void sample(MatrixXd &src)
	{
		uint32_t i;

		for (i=0; i < src.array().size(); i++)
		{
			src.array()(i) = (double)(src.array()(i) > Noise_Uniform(&m_noise));
		}
	}

	void probsLogistic(MatrixXd &src)
	{
		src.array() = (-src.array()).exp();
		src.array() += 1;
		src.array() = 1.0/src.array();
	}

	void probsLogistic(RowVectorXd &src)
	{
		src.array() = (-src.array()).exp();
		src.array() += 1;
		src.array() = 1.0/src.array();
	}

	void probsLogistic(MatrixXd &src, const MatrixXd &sigma)
	{
		src.array() /= (sigma.array() + EPSILON_SIGMA);
		src.array() = (-src.array()).exp();
		src.array() += 1;
		src.array() = 1.0/src.array();
	}

	void probsLogistic(RowVectorXd &src, const RowVectorXd &sigma)
	{
		src.array() /= (sigma.array() + EPSILON_SIGMA);
		src.array() = (-src.array()).exp();
		src.array() += 1;
		src.array() = 1.0/src.array();
	}

	void probsGaussian(MatrixXd &src, const MatrixXd &sigma)
	{
		src.array() = 1 - src.array();
		src.array() *= src.array();
		src.array() *= -0.5;

		MatrixXd var = sigma;
		var.array() += EPSILON_SIGMA;
		var.array() *= var.array();

		src.array() /= var.array();
		src.array() = src.array().exp();

		MatrixXd k = var;

		k.array() *= 2*3.14159265359;
		k.array() = k.array().sqrt();
		k.array() = 1.0/k.array();

		src.array() *= k.array();
	}

	void probsGaussian(RowVectorXd &src, const RowVectorXd &sigma)
	{
		src.array() = 1 - src.array();
		src.array() *= src.array();
		src.array() *= -0.5;

		RowVectorXd var = sigma;
		var.array() += EPSILON_SIGMA;
		var.array() *= var.array();

		src.array() /= var.array();
		src.array() = src.array().exp();

		RowVectorXd k = var;

		k.array() *= 2*3.14159265359;
		k.array() = k.array().sqrt();
		k.array() = 1.0/k.array();

		src.array() *= k.array();
	}

	void sampleGaussian(MatrixXd &src, const MatrixXd &sigma)
	{
		uint32_t i;

		for (i=0; i < src.array().size(); i++)
		{
			src.array()(i) = sigma(i)*Noise_Gaussian(&m_noise) + src.array()(i);
		}
	}

	RowVectorXd normalizeData(RowVectorXd const &src, RowVectorXd const &mu, RowVectorXd const &var)
	{
		// Remove mean
		RowVectorXd res = src - mu;
//		res.array() /= var.array() + EPSILON_SIGMA;

//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
		return res;
	}

	RowVectorXd calcMean(MatrixXd const &batch)
	{
		// Remove mean
		RowVectorXd res = batch.colwise().mean();

//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
		return res;

	}

	RowVectorXd calcSigma(MatrixXd const &batch)
	{
		MatrixXd x = batch.rowwise() - batch.colwise().mean();

		x.array() *= x.array();

		RowVectorXd res = x.colwise().mean().array().sqrt();

//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
		return res;
	}

	MatrixXd calcZ(MatrixXd &v, MatrixXd &h)
	{

		MatrixXd t1(v.rows(), m_w.getNumVisible());

		t1 = v - m_w.visibleBias().transpose().replicate(v.rows(), 1);
		t1.array() *= t1.array();
		t1.array() *= 0.5;

		t1 -= (h * m_w.weights().transpose());

		return t1;
	}

	void train(uint32_t numEpochs, double sigmaMin = 0.05)
	{
		uint32_t t, i;
		uint32_t epoch;
		uint32_t gibbs;

		size_t batchSize = m_batch.rows();

		double dProgress = 1.0/numEpochs;
		double kTrain = 1.0/batchSize;


		m_v.resize(batchSize, m_w.getNumVisible());
		m_h.resize(batchSize, m_w.getNumHidden());
		
		MatrixXd sumBiasV(1, m_w.getNumVisible());
		MatrixXd sumBiasH(1, m_w.getNumHidden());
		MatrixXd sumWeights(m_w.getNumVisible(), m_w.getNumHidden());

		MatrixXd deltaBiasV(MatrixXd::Zero(1, m_w.getNumVisible()));
		MatrixXd deltaBiasH(MatrixXd::Zero(1, m_w.getNumHidden()));
		MatrixXd deltaWeights(MatrixXd::Zero(m_w.getNumVisible(), m_w.getNumHidden()));

		MatrixXd diffErr(batchSize, m_w.getNumVisible());

		m_progress = 0;

		MatrixXd batch = m_batch;

		if (m_doNormalizeData)
		{
			RowVectorXd mean = calcMean(batch);
			for (i=0; i < batchSize; i++)
			{
				RowVectorXd x = batch.row(i);
				batch.row(i) = normalizeData(x, mean, m_variableSigma);
			}
		}

		for (epoch=0; epoch < numEpochs; epoch++)
		{
			double err;

			// Create hidden layer base on training data
			toHiddenBatch(m_h, batch);
			probsLogistic(m_h);

			if (!m_doRaoBlackwell)
			{
				sample(m_h);
			}

			// Update weights (positive phase)
			sumBiasV = batch.colwise().sum();
			if (!m_doSparse)
			{
				sumBiasH = m_h.colwise().sum();
			}
			sumWeights = batch.transpose() * m_h;
			diffErr = batch;

			for (gibbs=0; gibbs < m_numGibbs; gibbs++)
			{
				sample(m_h);

				// Create visible reconstruction (a fantasy...) given h
				toVisibleBatch(m_v, m_h);
				if (m_useVisibleGaussian)
				{
					if (!m_useProbsForHiddenReconstruction)
					{
						sampleGaussian(m_v, m_variableSigma.replicate(batchSize, 1));
					}
				}
				else
				{
					probsLogistic(m_v, m_variableSigma.replicate(batchSize, 1));
					if (!m_useProbsForHiddenReconstruction)
					{
						sample(m_v);
					}
				}

				// Create hidden representation given v
				toHiddenBatch(m_h, m_v);
				probsLogistic(m_h);
			}

			if (!m_doRaoBlackwell)
			{
				sample(m_h);
			}
			// Update weights (negative phase)
			sumBiasV -= m_v.colwise().sum();
			if (!m_doSparse)
			{
				sumBiasH -= m_h.colwise().sum();
			}
			sumWeights -= m_v.transpose() * m_h;
			diffErr -= m_v;

			deltaWeights = m_momentum*deltaWeights + m_muWeights*(kTrain*sumWeights - m_weightDecay*m_w.weights());
			m_w.weights() += deltaWeights;

			deltaBiasV = m_momentum*deltaBiasV + m_muWeights*kTrain*sumBiasV;
			m_w.visibleBias() += deltaBiasV;

			if (m_doSparse)
			{
				// Create hidden representation given v
				toHiddenBatch(m_h, m_v);
				probsLogistic(m_h);

				sumBiasH.fill(m_sparsity);
				sumBiasH -= m_h.colwise().mean();

				deltaBiasH = m_momentum*deltaBiasH + m_muSparsity*sumBiasH;

//				cout << "Mean(" << m_sparsity << ") = " << (double)sumBiasH.array().mean() << endl;
//				cout << sumBiasH << endl;
			}
			else
			{
				deltaBiasH = m_momentum*deltaBiasH + m_muWeights*kTrain*sumBiasH;
			}
			m_w.hiddenBias() += deltaBiasH;

			if (m_variableSigma[0] > sigmaMin)
			{
				m_variableSigma.array() *= m_sigmaDecay;
			}

			m_progress += dProgress;
			if (m_pListener)
			{
				m_pListener->onEpochTrained(*this);
			}

			diffErr.array() *= diffErr.array();
			err = diffErr.colwise().sum().sum();

			cout << "err =" << endl;
			cout << err << endl;

		} // Number of epochs
	}

	double getProgress() const
	{
		return m_progress;
	}

	double getEnergy(const VectorXd& visible, const VectorXd& hidden)
	{
		double energy;
		double sigma = m_variableSigma.array().mean();

		energy = m_w.visibleBias() * visible;
		energy += m_w.hiddenBias() * hidden;
		energy += visible.transpose() * m_w.weights() * hidden;
		
		return -energy/(sigma*sigma);
	}

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

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

		if (m_useVisibleGaussian)
		{
//			probsGaussian(v, m_sigmas);
		}
		else
		{
			probsLogistic(v, m_variableSigma);
		}
	}

	void setConstantSigma(double value)
	{	
		 m_constantSigma = value;
		 m_variableSigma.fill(m_constantSigma);
	}

	double getConstantSigma()
	{
		 return m_constantSigma;
	}

	RowVectorXd& getVariableSigma()
	{
		 return m_variableSigma;
	}

	void setSigmaDecay(double value)
	{
		m_sigmaDecay = value;
	}

	void setWeightDecay(double value)
	{
		m_weightDecay = value;
	}

	void setLambda(double value)
	{
		m_lambda = value;
	}

	void setSparsity(double value)
	{
		m_sparsity = value;
	}

	void setUseVisibleGaussian(bool flag)
	{
		m_useVisibleGaussian = flag;
	}

	void setDoRaoBlackwell(bool flag)
	{
		m_doRaoBlackwell = flag;
	}

	void setUseProbsForHiddenReconstruction(bool flag)
	{
		m_useProbsForHiddenReconstruction = flag;
	}

	void setDoSparse(bool flag)
	{
		m_doSparse = flag;
	}

	void setNormalizeData(bool flag)
	{
		m_doNormalizeData = flag;
	}

	void setDoLearnVariance(bool flag)
	{
		m_doLearnVariance = flag;
		if (m_doLearnVariance && m_batch.rows())
		{
			m_variableSigma = calcSigma(m_batch);
		}
		else
		{
			 m_variableSigma.fill(m_constantSigma);
		}
	}

	void setNumGibbs(uint32_t value)
	{
		m_numGibbs = value;
	}

	uint32_t getNumGibbs()
	{
		return m_numGibbs;
	}

	void setMuWeights(double value)
	{
		m_muWeights = value;
	}

	void setMuSparsity(double value)
	{
		m_muSparsity = value;
	}

	void setMomentum(double value)
	{
		m_momentum = value;
	}

	MatrixXd const& getHiddenBatch()
	{
		return m_h;
	}
	
	MatrixXd const& getVisibleBatch()
	{
		return m_v;
	}

	MatrixXd const& getBatch()
	{
		return m_batch;
	}
	
private:
	Weights &m_w;
	MatrixXd const &m_batch;
	MatrixXd m_v;
	MatrixXd m_h;
	RowVectorXd m_variableSigma;
	double m_constantSigma;
	RbmListener *m_pListener;
	noise_gen_t m_noise;
	double m_progress;
	double m_sigmaDecay;
	double m_weightDecay;
	double m_lambda;
	double m_sparsity;
	double m_muWeights;
	double m_muSparsity;
	double m_momentum;
	bool m_useVisibleGaussian;
	bool m_doRaoBlackwell;
	bool m_useProbsForHiddenReconstruction;
	bool m_doSparse;
	bool m_doNormalizeData;
	bool m_doLearnVariance;
	uint32_t m_numGibbs;

	void toHiddenBatch(MatrixXd &h, MatrixXd const &v)
	{
		h = v * m_w.weights();
		h += m_w.hiddenBias().replicate(m_batch.rows(), 1);
	}

	void toVisibleBatch(MatrixXd &v, MatrixXd const &h)
	{
		v = h * m_w.weights().transpose();
		v += m_w.visibleBias().replicate(m_batch.rows(), 1);
	}
};




#endif /* RBM_HPP_ */