/*
 * 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", (int)numGibbs).asUInt();
			miniBatchSize = params.get("miniBatchSize", (int)miniBatchSize).asUInt();
			numEpochs = params.get("numEpochs", (int)numEpochs).asUInt();
		}
		
		double weightDecay;
		double learningRate;
		double momentum;
		bool doRaoBlackwell;
		bool gibbsDoSampleVisible;
		bool gibbsDoSampleHidden;
		bool doSampleBatch;
		size_t numGibbs;
		size_t miniBatchSize;
		size_t numEpochs;
	};

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

	class IListener
	{
	public:

		IListener() {}
		virtual ~IListener() {}
		virtual bool onProgress(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 train(arma::mat const &batch, IListener *pListener=nullptr);
	
	arma::mat toHiddenState(const arma::mat &visible) const;
	arma::mat toVisibleState(const arma::mat &hidden) const;
	arma::mat toHiddenProbs(const arma::mat &visible) const;
	arma::mat toVisibleProbs(const arma::mat &hidden) const;
	const arma::mat& w() 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;
	}
private:
	Params m_params;
	arma::mat sample(arma::mat const &src);
	static arma::mat probsLogistic(arma::mat const &src);
	void uniform(arma::mat &srcDst, double stdDev=1.0, double mu=0.5);

protected:
	arma::mat m_w;
	arma::mat m_bh;
	arma::mat m_bv;

private:
	noise_gen_t m_noise;
};

#endif /* RBM_HPP */