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

/* 
 * File:   Rbm.hpp
 * Author: jens
 *
 * Created on 21. Oktober 2019, 21:28
 */

#ifndef RBM_HPP
#define RBM_HPP

#include <streambuf>

#include <armadillo>
#include <jsoncpp/json/json.h>
#include "noise.h"

class Rbm
{
public:
	
	struct Params
	{
		Params()
		: learningRate(0.1)
		, weightDecay(0.0)
		, momentum(0.5)
		, doRaoBlackwell(true)
		, gibbsDoSampleVisible(false)
		, gibbsDoSampleHidden(true)
		, doSampleBatch(false)
		, numGibbs(1)
		, miniBatchSize(100)
		, numEpochs(1000)
		{
		}

		Json::Value toJson() const
		{
			std::cout << "Exporting Rbm::Params" << std::endl;
			Json::Value params;
			params["weightDecay"] = weightDecay;
			params["learningRate"] = learningRate;
			params["momentum"] = momentum;
			params["doRaoBlackwell"] = (int)doRaoBlackwell;
			params["gibbsDoSampleVisible"] = (int)gibbsDoSampleVisible;
			params["gibbsDoSampleHidden"] = (int)gibbsDoSampleHidden;
			params["doSampleBatch"] = (int)doSampleBatch;
			params["numGibbs"] = (int)numGibbs;
			params["miniBatchSize"] = (int)miniBatchSize;
			params["numEpochs"] = (int)numEpochs;

			return params;
		}

		void fromJson(Json::Value params)
		{
			std::cout << "Importing Rbm::Params" << std::endl;
			weightDecay = params.get("weightDecay", weightDecay).asDouble();
			learningRate = params.get("learningRate", learningRate).asDouble();
			momentum = params.get("momentum", momentum).asDouble();
			doRaoBlackwell = params.get("doRaoBlackwell", doRaoBlackwell) == 1;
			gibbsDoSampleVisible = params.get("gibbsDoSampleVisible", gibbsDoSampleVisible) == 1;
			gibbsDoSampleHidden = params.get("gibbsDoSampleHidden", gibbsDoSampleHidden) == 1;
			doSampleBatch = params.get("doSampleBatch", doSampleBatch) == 1;
			numGibbs = params.get("numGibbs", numGibbs).asUInt();
			miniBatchSize = params.get("miniBatchSize", miniBatchSize).asUInt();
			numEpochs = params.get("numEpochs", numEpochs).asUInt();
		}
		
		double weightDecay;
		double learningRate;
		double momentum;
		bool doRaoBlackwell;
		bool gibbsDoSampleVisible;
		bool gibbsDoSampleHidden;
		bool doSampleBatch;
		int numGibbs;
		int miniBatchSize;
		int numEpochs;
	};

	struct Status
	{
		Status()
		: progress(0)
		, err(-1.0)
		, err_total(-1.0)
		, L1(-1.0)
		, L2(-1.0)
		{
		}
		int progress;
		double err;
		double err_total;
		double L1;
		double L2;
	};

	class IListener
	{
	public:

		IListener() {}
		virtual ~IListener() {}
		virtual bool onProgress(Rbm *pRbm, const Status &status)
		{
			return true;
		}
	};

	Rbm(size_t numVisible, size_t numHidden);
	Rbm(const Rbm& orig);
	virtual ~Rbm();

	void weightsInit(double stddev, double mu=0.0);
	void weightsAssign(const arma::mat &w, const arma::mat &bhv, const arma::mat &bv)
	{
		m_whv.submat(0, 0, w.n_rows-1, w.n_cols-1) = w;
		m_bhv.submat(0, 0, bhv.n_rows-1, bhv.n_cols-1) = bhv;
		m_bv.submat(0, 0, bv.n_rows-1, bv.n_cols-1) = bv;
	}
	
	void setBatch(arma::mat const &batch)
	{
		m_batch = batch;
	}
	
	void train(IListener *pListener=nullptr);
	
	static arma::mat normalize(const arma::mat &hidden);
	const arma::mat& whv() const;
	const arma::mat& bv() const;
	const arma::mat& bh() const;
	
	Json::Value toJson() const;
	void fromJson(Json::Value params);
	
	Params& params()
	{
		return m_params;
	}
	static arma::mat prob(arma::mat const &src);
     
	arma::mat toHiddenProbs(const arma::mat &visible) const
	{
		return Rbm::prob(v_to_h(visible));
	}

	arma::mat toVisibleProbs(const arma::mat &hidden) const
	{
		return Rbm::prob(h_to_v(hidden));
	}
	
	size_t numHidden() const
	{
		return m_bhv.size();
	}

	size_t numVisible() const
	{
		return m_bv.size();
	}
	
	arma::mat v_to_h(const arma::mat &visible) const;
	arma::mat h_to_v(const arma::mat &hidden) const;

private:
	Params m_params;
	arma::mat sample(arma::mat const &src);
	void contrastiveDivergence(arma::mat const &v_states, arma::mat &dwhv, arma::mat &dbhv, arma::mat &dbv);
	void uniform(arma::mat &srcDst, double stdDev=1.0, double mu=0.5);
	void gibbs(arma::mat &hv_states, arma::mat &v_states);
	
protected:
	arma::mat m_bhv;
	arma::mat m_bv;

private:
	noise_gen_t m_noise;
	arma::mat m_whv;
	arma::mat m_batch;
	
};

#endif /* RBM_HPP */