/*
 * 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)
	{
		Noise_Init(&m_noise, 0x32727155);
	}

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

	void weightsUpdate(VisibleLayer &v, HiddenLayer &h, double mu)
	{
		MatrixXd &w = (MatrixXd&)m_w.weights();

		// Update weights
		w += mu*(v.states() * h.states().transpose());
	}

	void visibleBiasUpdate(VisibleLayer &v, double mu)
	{
		m_w.visibleBias().array() += mu*v.states().array();
	}

	void hiddenBiasUpdate(HiddenLayer &h, double mu)
	{
		m_w.hiddenBias().array() += mu*h.states().array();
	}
//#define RBM_SPARSE
	void train(LayerArray<VisibleLayer> &vt, uint32_t numEpochs, double mu, uint32_t numGibbs = 1, bool useExpectations = false, bool doRaoBlackwell = false, bool useProbsForHiddenReconstruction = false, bool doRobbinsMonro = false)
	{
		uint32_t t, i;
		uint32_t epoch;
		uint32_t gibbs;
		VisibleLayer v(m_w.getNumVisible());
		HiddenLayer h(m_w.getNumHidden());
		HiddenLayer *pH;
		LayerArray<HiddenLayer> ht(vt.getSize(), m_w.getNumHidden());

		Weights w = m_w;

		double dProgress = 1.0/numEpochs;
		m_progress = 0;

#ifdef RBM_SPARSE
		const double lambda = 0.05;
		const double variance = 0.4;
		const double penalty = 0.05;
#else
		const double lambda = 1.0;
		const double variance = 1.0;
		const double penalty = 0.0;
#endif

		if (useExpectations)
		{
			mu /= vt.getSize();
		}

		if (doRobbinsMonro)
		{
			for (i=0; i < vt.getSize(); i++)
			{
				// Create hidden layer base on training data
				ht[i].probsUpdateLogistic(vt[i], w, lambda, variance);
			}
		}

		for (epoch=0; epoch < numEpochs; epoch++)
		{
			for (i=0; i < vt.getSize(); i++)
			{
				t = (uint32_t)(0.5 + (vt.getSize()-1)*Noise_Uniform(&m_noise, 0.5));
				h.probsUpdateLogistic(vt[t], w, lambda, variance);

				// Create hidden layer base on training data
				if (doRobbinsMonro)
				{
					pH = &ht[t];
				}
				else
				{
					pH = &h;
				}

				// Update weights (positive phase)
				if (doRaoBlackwell)
				{
					pH->states() = pH->probs();
				}
				else
				{
					pH->statesUpdateStochastic();
				}
				weightsUpdate(vt[t], h, +mu);
				visibleBiasUpdate(vt[t], +mu);
				hiddenBiasUpdate(h, +mu);

				for (gibbs=0; gibbs < numGibbs; gibbs++)
				{
					pH->statesUpdateStochastic();

					// Create visible reconstruction (a fantasy...)
#ifdef RBM_SPARSE
					v.probsUpdateGaussian(*pH, w, lambda, variance);
#else
					v.probsUpdateLogistic(*pH, w);
#endif

					if (useProbsForHiddenReconstruction)
					{
						v.states() = v.probs();
					}
					else
					{
						v.statesUpdateStochastic();
					}

					// Create hidden reconstruction
					pH->probsUpdateLogistic(v, w, lambda, variance);
				}
				// Update weights (negative phase)
				if (doRaoBlackwell)
				{
					pH->states() = pH->probs();
				}
				else
				{
					pH->statesUpdateStochastic();
				}
				weightsUpdate(v, *pH, -mu);
				visibleBiasUpdate(v, -mu);
				hiddenBiasUpdate(*pH, -mu);

				if (!useExpectations)
				{
					w = m_w;
				}

			}
#ifdef RBM_SPARSE
			{
				HiddenLayer th(m_w.getNumHidden());
				VectorXd m(th.states());
				m.fill(0);
				for (i=0; i < vt.getSize(); i++)
				{
					m += expectHidden(vt[i].states(), lambda, variance, 10);
				}
				m.array() = penalty - m.array();
				m *= 1.0/vt.getSize();
				th.states() = m;
				hiddenBiasUpdate(th, -mu);
			}
#endif

			if (useExpectations)
			{
				w = m_w;
			}

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

	double getProgress() const
	{
		return m_progress;
	}

	double getEnergy(VisibleLayer &v, HiddenLayer &h)
	{
		uint32_t i, j;
		double energy;

		energy = -v.getEnergy(m_w) - h.getEnergy(m_w);

		for (i=0; i < h.getNumUnits(); i++)
		{
			for (j=0; j < v.getNumUnits(); j++)
			{
//				energy -= v.getStates()[j] * h.getStates()[i] * m_w.getWeights()[i][j];
			}
		}

		// ToDo: make this correct
//		energy -= (v.states().transpose() * h.states()); // * m_w.weights();
		return energy;
	}

	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);
//			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), h[i]));
			}
			for (j=0; j < vts.getSize(); j++)
			{
				p = exp(-getEnergy(vts.getAt(j), h[i]))/z;
				cout << p << endl;
			}
			cout << endl;
		}
		cout << endl;

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

		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, double lambda = 1.0, double variance = 1.0)
	{
		HiddenLayer th(m_w.getNumHidden());
		VisibleLayer tv(m_w.getNumVisible(), (const VectorXd*)&visible);

		th.probsUpdateLogistic(tv, m_w, lambda, variance);

		return th.probs();
	}

	VectorXd toVisible(const VectorXd& hidden, double lambda = 1.0, double variance = 1.0)
	{
		HiddenLayer th(m_w.getNumHidden(), (const VectorXd*)&hidden);
		VisibleLayer tv(m_w.getNumVisible());

		tv.probsUpdateLogistic(th, m_w, lambda, variance);

		return tv.probs();
	}

	VectorXd expectHidden(VectorXd visible, uint32_t numIter, double lambda = 1.0, double variance = 1.0)
	{
      	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, lambda, variance);
    		v.probsUpdateGaussian(h, (Weights&)m_w, lambda, variance);
    	}

		return h.probs();
	}

   VectorXd expectVisible(VectorXd visible, uint32_t numIter, double lambda = 1.0, double variance = 1.0)
	{
      	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, lambda, variance);
    		v.probsUpdateLogistic(h, (Weights&)m_w, lambda, variance);
    	}

		return v.probs();
	}

private:
	Weights &m_w;
	RbmListener *m_pListener;
	noise_gen_t m_noise;
	double m_progress;

};




#endif /* RBM_HPP_ */