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

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_tv(weights.getNumVisible())
	, m_th(weights.getNumHidden())
	, m_progress(0)
	{
		Noise_Init(&m_noise, 0x32727155);
	}

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

	void weightsUpdate(VisibleLayer &v, HiddenLayer &h, double mu)
	{
		uint32_t i, j;
		double dw;

		double **ppW = m_w.getWeights();
		const double *pH = h.getStates();
		const double *pV = v.getStates();

		// Update weights
		for (i=0; i < m_w.getNumHidden(); i++)
		{
			for (j=0; j < m_w.getNumVisible(); j++)
			{
				dw = pV[j] * pH[i];
				ppW[i][j] += mu*dw;
			}
		}
		double *pBias = m_w.getBiasVisible();
		for (i=0; i < m_w.getNumVisible(); i++)
		{
			dw = pV[i];
			pBias[i] += mu*dw;
		}

		pBias = m_w.getBiasHidden();
		for (i=0; i < m_w.getNumHidden(); i++)
		{
			dw = pH[i];
			pBias[i] += mu*dw;
		}
	}

	void train(LayerArray<VisibleLayer> &vts, uint32_t numEpochs, double mu, uint32_t numGibbs = 1, bool useExpectations = false, bool doRaoBlackwell = false, bool doRobinsMonro = false)
	{
		uint32_t t;
		uint32_t epoch;
		uint32_t gibbs;
		VisibleLayer v(m_w.getNumVisible());
		HiddenLayer h(m_w.getNumHidden());
		Weights w = m_w;

		double dProgress = 1.0/numEpochs;
		m_progress = 0;

		for (epoch=0; epoch < numEpochs; epoch++)
		{
			for (t=0; t < vts.getSize(); t++)
			{
				// Create hidden layer base on training data
				h.probsUpdate(vts[t], w);

				// Update weights (positive phase)
				if (doRaoBlackwell)
				{
					h.statesAssignfromProbs();
				}
				else
				{
					h.statesUpdateStochastic();
				}
				weightsUpdate(vts[t], h, +mu/vts.getSize());

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

					// Create visible reconstruction (a fantasy...)
					v.probsUpdate(h, w);
					v.statesUpdateStochastic();

					// Create hidden reconstruction
					h.probsUpdate(v, w);
				}
				// Update weights (negative phase)
				if (doRaoBlackwell)
				{
					h.statesAssignfromProbs();
				}
				else
				{
					h.statesUpdateStochastic();
				}
				weightsUpdate(v, h, -mu/vts.getSize());

				if (!useExpectations)
				{
					w = m_w;
				}

			}
			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];
			}
		}
		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].probsUpdate(vts.getAt(j), m_w);
//			h[j].statesAssignfromProbs();
			h[j].statesUpdateStochastic();
		}

		printf("pi(t) = (pi^, v>)\n");
		for (i=0; i < m_w.getNumHidden(); i++)
		{
			for (j=0; j < vts.getSize(); j++)
			{
				p = h[j].getProbs()[i];
				printf("%3.6f ", p);
			}
			printf("\n");
		}
		printf("\n");

		printf("si(t) = (si^, v>)\n");
		for (i=0; i < m_w.getNumHidden(); i++)
		{
			for (j=0; j < vts.getSize(); j++)
			{
				p = h[j].getStates()[i];
				printf("%3.6f ", p);
			}
			printf("\n");
		}
		printf("\n");

		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;
				printf("%3.6f ", p);
			}
			printf("\n");
		}
		printf("\n");

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

		printf("A fantasy... (v^, t>)\n");
		for (i=0; i < m_w.getNumVisible(); i++)
		{
			for (j=0; j < vts.getSize(); j++)
			{
				p = vts.getAt(j).getProbs()[i];
				printf("%3.6f ", p);
			}
			printf("\n");
		}

		delete [] h;
	}

	const double* toHidden(const double *pVisible)
	{
		double p;
		uint32_t i;

		VisibleLayer tv(m_w.getNumVisible(), pVisible);

		m_th.probsUpdate(tv, m_w);
		m_th.statesAssignfromProbs();
//		m_th.statesUpdateStochastic();

#if 0
		printf("pi(t) = (pi^, v>)\n");
		for (i=0; i < m_w.getNumHidden(); i++)
		{
			p = m_th.getProbs()[i];
			printf("%3.6f\n", p);
		}
		printf("\n");
#endif
		return m_th.getStates();
	}

	const double* toVisible(const double *pHidden)
	{
		double p;
		uint32_t i;

		HiddenLayer th(m_w.getNumHidden(), pHidden);

		m_tv.probsUpdate(th, m_w);
		m_tv.statesAssignfromProbs();
//		m_tv.statesUpdateStochastic();

#if 0
		printf("pi(t) = (pi^, v>)\n");
		for (i=0; i < m_w.getNumVisible(); i++)
		{
			p = m_tv.getProbs()[i];
			printf("%3.6f\n", p);
		}
		printf("\n");
#endif
		return m_tv.getStates();
	}

private:
	Weights &m_w;
	RbmListener *m_pListener;
	VisibleLayer m_tv;
	HiddenLayer m_th;
	noise_gen_t m_noise;
	double m_progress;

};




#endif /* RBM_HPP_ */