/*
 * 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 <cmath>
#include <Eigen/Dense>

using namespace Eigen;

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

class Rbm;

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

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

class Rbm
{
public:
	Rbm(Weights &weights, RbmListener *pListener = nullptr)
	: m_w(weights)
	, m_pListener(pListener)
	, m_progress(0)
	, m_sigma(1.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_doCancel(false)
	, m_useVisibleGaussian(false)
	, m_doRaoBlackwell(false)
	, m_useProbsForHiddenReconstruction(false)
	, m_doSparse(false)
	, m_doNormalizeData(false)
	, m_numGibbs(1)
	{
		Noise_Init(&m_noise, 0x32727155);
	}

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

	void train(const LayerArray<VisibleLayer> &batch, uint32_t numEpochs, double sigmaMin = 0.05)
	{
		uint32_t t, i;
		uint32_t epoch;
		uint32_t gibbs;
		double sigma;
		VisibleLayer v(m_w.getNumVisible());
		HiddenLayer h(m_w.getNumHidden());

		VectorXd sumBiasV(m_w.getNumVisible());
		VectorXd deltaBiasV(m_w.getNumVisible());

		VectorXd sumBiasH(m_w.getNumHidden());
		VectorXd deltaBiasH(m_w.getNumHidden());

		MatrixXd sumWeights(m_w.getNumVisible(), m_w.getNumHidden());
		MatrixXd deltaWeights(m_w.getNumVisible(), m_w.getNumHidden());

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

		const LayerArray<VisibleLayer> &vt = batch;

		sigma = m_sigma;

		double dProgress = 1.0/numEpochs;
		double kTrain = 1.0/vt.getSize();

		m_progress = 0;

		deltaWeights.fill(0);
		deltaBiasV.fill(0);
		deltaBiasH.fill(0);
		m_doCancel = false;
		for (epoch=0; epoch < numEpochs; epoch++)
		{
			double err = 0;

			if (m_doCancel)
			{
				m_doCancel = false;
				break;
			}

			sumWeights.fill(0);
			sumBiasV.fill(0);
			sumBiasH.fill(0);
			for (i=0; i < vt.getSize(); i++)
			{
				t = i;
				h.probsUpdateLogistic(vt[t], m_w, m_lambda, sigma);

				// Create hidden layer base on training data
				if (m_doRaoBlackwell)
				{
					h.states() = h.probs();
				}
				else
				{
					h.statesUpdateStochastic();
				}

				// Update weights (positive phase)
				sumWeights += vt[t].states() * h.states().transpose();
				sumBiasV += vt[t].states();
				sumBiasH += h.states();

				diffErr = vt[t].states();

				for (gibbs=0; gibbs < m_numGibbs; gibbs++)
				{
					h.statesUpdateStochastic();

					// Create visible reconstruction (a fantasy...)
					if (m_useProbsForHiddenReconstruction)
					{
						if (m_useVisibleGaussian)
						{
							v.probsUpdateGaussian(h, m_w, m_lambda, sigma);
						}
						else
						{
							v.probsUpdateLogistic(h, m_w, m_lambda, sigma);
						}
						v.states() = v.probs();
					}
					else
					{
						if (m_useVisibleGaussian)
						{
							v.sampleGaussian(h, m_w, m_lambda, sigma);
						}
						else
						{
							v.probsUpdateLogistic(h, m_w, m_lambda, sigma);
							v.statesUpdateStochastic();
						}
					}
					// Create hidden reconstruction
					h.probsUpdateLogistic(v, m_w, m_lambda, sigma);
				}

				// Update weights (negative phase)
				if (m_doRaoBlackwell)
				{
					h.states() = h.probs();
				}
				else
				{
					h.statesUpdateStochastic();
				}
				sumWeights -= v.states() * h.states().transpose();
				sumBiasV -= v.states();
				sumBiasH -= h.states();
				diffErr -= v.states();
				diffErr.array() *= diffErr.array();
				err += diffErr.sum();

			} // TrainingSize

			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)
			{
				HiddenLayer th(m_w.getNumHidden());
				VectorXd m(m_w.getNumHidden());
				m.fill(0);

				for (i=0; i < vt.getSize(); i++)
				{
					th.probsUpdateLogistic(vt[i], m_w, m_lambda, sigma);
					m += th.probs();
				}
				m /= i;
				sumBiasH = m_sparsity - m.array();
				deltaBiasH = m_momentum*deltaBiasH + m_muSparsity*sumBiasH;

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

			m_w.hiddenBias() += deltaBiasH;

			if (sigma > sigmaMin)
			{
				sigma *= m_sigmaDecay;
			}

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

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

		} // Number of epochs
	}

	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, double lambda, double sigma)
	{
		double var = sigma*sigma;

		src.array() *= -lambda/var;
		src.array() = src.array().exp();
		src.array() += 1;
		src.array() = 1.0/src.array();
	}

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

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

	void normalizeData(MatrixXd &src, double mu, double sigma)
	{
		uint32_t i;
		uint32_t size = src.rows();
		double mean;
		double stdDev;

		for (i=0; i < size; i++)
		{
			mean = src.row(i).array().mean();
			src.row(i).array() -= mean;
			src.row(i).array() += mu;
		}
		for (i=0; i < size; i++)
		{
			src.row(i).array() *= src.row(i).array();
		}
		for (i=0; i < size; i++)
		{
			stdDev = sqrt(src.row(i).array().mean());
			src.row(i).array() /= stdDev;
			src.row(i).array() *= sigma;
		}
	}

	void train2(const LayerArray<VisibleLayer> &vt, uint32_t numEpochs, uint32_t batchSize, double sigmaMin = 0.05)
	{
		uint32_t t, i;
		uint32_t epoch;
		uint32_t gibbs;
		double sigma = m_sigma;
		double dProgress = 1.0/numEpochs;
		double kTrain = 1.0/vt.getSize();

//		if (batchSize > vt.getSize())
		batchSize = vt.getSize();

		MatrixXd v(batchSize, m_w.getNumVisible());
		MatrixXd h(batchSize, m_w.getNumHidden());
		MatrixXd batch(batchSize, m_w.getNumVisible());

		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;
		m_doCancel = false;

		for (i=0; i < batchSize; i++)
		{
//			t = (uint32_t)(0.5 + (vt.getSize()-1)*Noise_Uniform(&m_noise));
			batch.row(i) = vt[i].states();

		}
		if (m_doNormalizeData)
		{
			normalizeData(batch, 0.0, m_sigma);
		}

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

			v = batch;
			if (m_doCancel)
			{
				m_doCancel = false;
				break;
			}

			// Create hidden layer base on training data
			h = v * m_w.weights();
			h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
			probsLogistic(h, m_lambda, sigma);

			if (!m_doRaoBlackwell)
			{
				sample(h);
			}
			// Update weights (positive phase)
			sumBiasV = v.colwise().sum();
			if (!m_doSparse)
			{
				sumBiasH = h.colwise().sum();
			}
			sumWeights = v.transpose() * h;
			diffErr = v;

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

				// Create visible reconstruction (a fantasy...)
				v = h * m_w.weights().transpose();
				v += m_w.visibleBias().transpose().replicate(batchSize, 1);

				if (m_useVisibleGaussian)
				{
					sampleGaussian(v, m_lambda, sigma);
				}
				else
				{
					probsLogistic(v, m_lambda, sigma);
					if (!m_useProbsForHiddenReconstruction)
					{
						sample(v);
					}
				}

				// Create hidden reconstruction
				h = v * m_w.weights();
				h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
				probsLogistic(h, m_lambda, sigma);
			}

			if (!m_doRaoBlackwell)
			{
				sample(h);
			}
			// Update weights (negative phase)
			sumBiasV -= v.colwise().sum();
			if (!m_doSparse)
			{
				sumBiasH -= h.colwise().sum();
			}
			sumWeights -= v.transpose() * h;
			diffErr -= 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)
			{
				h = v * m_w.weights();
				h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
				probsLogistic(h, m_lambda, sigma);

				sumBiasH.fill(m_sparsity);
				sumBiasH -= 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 (sigma > sigmaMin)
			{
				sigma *= 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;

		energy = m_w.visibleBias().transpose() * visible;
		energy += m_w.hiddenBias().transpose() * hidden;
		energy += visible.transpose() * m_w.weights() * hidden;

		return -energy/(m_sigma*m_sigma);
	}

	void prob(LayerArray<VisibleLayer> &vts)
	{
		uint32_t i, j;
		double z;
		double p;

		HiddenLayer *h = new HiddenLayer[vts.getSize()];

		// Create hidden layer activations based on training data
		for (j=0; j < vts.getSize(); j++)
		{
			h[j].setNumUnits(m_w.getNumHidden());
			h[j].probsUpdateLogistic(vts.getAt(j), m_w, m_lambda, m_sigma);
//			h[j].statesAssignfromProbs();
			h[j].statesUpdateStochastic();
		}

		printf("pi(t) = (pi^, v>)\n");
		for (j=0; j < vts.getSize(); j++)
		{
			cout << h[j].probs() << endl;
		}
		cout << endl;

		printf("si(t) = (si^, v>)\n");
		for (j=0; j < vts.getSize(); j++)
		{
			cout << h[j].states() << endl;
		}
		cout << endl;

		printf("p(v) = (t^, v>)\n");
		for (i=0; i < vts.getSize(); i++)
		{
			z = 0;
			for (j=0; j < vts.getSize(); j++)
			{
				z += exp(-getEnergy(vts.getAt(j).states(), h[i].states()));
			}
			for (j=0; j < vts.getSize(); j++)
			{
				p = exp(-getEnergy(vts.getAt(j).states(), h[i].states()))/z;
				cout << p << endl;
			}
			cout << endl;
		}
		cout << endl;

		// Reconstruct
		for (i=0; i < vts.getSize(); i++)
		{
			vts.getAt(i).probsUpdateLogistic(h[i], m_w, m_lambda, m_sigma);
		}

		printf("A fantasy... (v^, t>)\n");
		for (j=0; j < vts.getSize(); j++)
		{
			cout << vts.getAt(j).probs() << endl;
		}

		delete [] h;
	}

	VectorXd toHidden(const VectorXd& visible)
	{
		HiddenLayer th(m_w.getNumHidden());
		VisibleLayer tv(m_w.getNumVisible(), (const VectorXd*)&visible);

		th.probsUpdateLogistic(tv, m_w, m_lambda, m_sigma);

		return th.probs();
	}

	VectorXd toVisible(const VectorXd& hidden)
	{
		HiddenLayer th(m_w.getNumHidden(), (const VectorXd*)&hidden);
		VisibleLayer tv(m_w.getNumVisible());
		if (m_useVisibleGaussian)
		{
			tv.probsUpdateGaussian(th, m_w, m_lambda, m_sigma);
		}
		else
		{
			tv.probsUpdateLogistic(th, m_w, m_lambda, m_sigma);
		}
		return tv.probs();
	}

	VectorXd expectHidden(VectorXd visible, uint32_t numIter)
	{
      	uint32_t i;
		VisibleLayer v(m_w.getNumVisible(), (const VectorXd*)&visible);
		HiddenLayer h(m_w.getNumHidden());

    	for (i=0; i < numIter; i++)
    	{
    		h.probsUpdateLogistic(v, (Weights&)m_w, m_lambda, m_sigma);
    		if (m_useVisibleGaussian)
    		{
    			v.probsUpdateGaussian(h, (Weights&)m_w, m_lambda, m_sigma);
    		}
    		else
    		{
    			v.probsUpdateLogistic(h, (Weights&)m_w, m_lambda, m_sigma);
    		}
    	}

		return h.probs();
	}

   VectorXd expectVisible(VectorXd visible, uint32_t numIter)
	{
      	uint32_t i;
		VisibleLayer v(m_w.getNumVisible(), (const VectorXd*)&visible);
		HiddenLayer h(m_w.getNumHidden());

    	for (i=0; i < numIter; i++)
    	{
    		h.probsUpdateLogistic(v, (Weights&)m_w, m_lambda, m_sigma);
    		if (m_useVisibleGaussian)
    		{
    			v.probsUpdateGaussian(h, (Weights&)m_w, m_lambda, m_sigma);
    		}
    		else
    		{
    			v.probsUpdateLogistic(h, (Weights&)m_w, m_lambda, m_sigma);
    		}
    	}

		return v.probs();
	}

   void setSigma(double value)
   {
	   m_sigma = value;
   }

   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 setNumGibbs(uint32_t value)
   {
	   m_numGibbs = value;
   }

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

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

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

   void cancel()
   {
	   m_doCancel = true;
//	   while(m_doCancel);
   }

private:
	Weights &m_w;
	RbmListener *m_pListener;
	noise_gen_t m_noise;
	double m_progress;
	double m_sigma;
	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;
	volatile bool m_doCancel;
	uint32_t m_numGibbs;

};




#endif /* RBM_HPP_ */