/*
 * 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:   Stack.cpp
 * Author: jens
 * 
 * Created on 25. Oktober 2019, 18:26
 */
#include <cassert>
#include "Stack.hpp"

using namespace std;

Stack::Stack(const std::string &name)
: m_name(name)
, m_pLayers(nullptr)
{
}

Stack::Stack(const Stack& orig)
: m_name(orig.m_name)
, m_pLayers(orig.m_pLayers)
{
}

Stack::~Stack()
{ 
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		Layer *pNextLayer = pLayer->next;
		delete pLayer;
		pLayer = pNextLayer;
	}
}

size_t Stack::numLayers()
{
	size_t count = 0;
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		count++;
		pLayer = pLayer->next;
	}
	return count;
}

void Stack::addLayer(Layer *pOtherLayer)
{
	if (!m_pLayers)
	{
		m_pLayers = pOtherLayer;
		pOtherLayer->prev = nullptr;
	}
	else
	{
		Layer *pLayer = m_pLayers;
		while(pLayer->next) 
		{
			pLayer = pLayer->next;
		}
		pLayer->next = pOtherLayer;
		pOtherLayer->prev = pLayer;
	}
}

void Stack::delLayer(Layer* pLayer)
{
	assert(!"Stack::delLayer: Not implemented!");
}

Layer* Stack::getLayer(size_t layerId) const
{
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		if (pLayer->id() == layerId)
		{
			return pLayer;
		}
		pLayer = pLayer->next;
	}
	return nullptr;
	
}

bool Stack::load(LayerConstructor *pLayerConstructor)
{
	std::cout << "Importing Project " << m_name << std::endl;
	ifstream ifs(m_name + string(".prj"));
	
	Json::Reader reader;
	Json::Value project;
	reader.parse(ifs, project);
	
	const string &name = project["stack"]["name"].asString();
	Json::Value &layers = project["stack"]["layers"];
			
	for (int i=0; i < layers.size(); i++)
	{
		Json::Value &layer = layers[i];
		
		string layername = layer["name"].asString();
		int numVisibleX = layer["numVisibleX"].asInt();
		int numVisibleY = layer["numVisibleY"].asInt();
		int numHidden = layer["numHidden"].asInt();
		
		Layer *pLayer = nullptr;
		if (!pLayerConstructor)
		{
			pLayer = new Layer(layername, i, numVisibleX, numVisibleY, numHidden);
		}
		else
		{
			pLayer = pLayerConstructor->onConstruct(layername, i, numVisibleX, numVisibleY, numHidden);
		}
		
		assert(pLayer != nullptr);
		
		pLayer->fromJson(layer["rbm"]);
		addLayer(pLayer);
	}
	
	return true;
}

bool Stack::save()
{
	std::cout << "Exporting Project " << m_name << std::endl;
	ofstream ofs(m_name + string(".prj"));
	
	Json::StyledWriter writer;
	Json::Value project;
	project["stack"]["name"] = m_name;
	
	Json::Value layers(Json::arrayValue);
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		layers.append(pLayer->toJson());
		pLayer = pLayer->next;
	}
	project["stack"]["layers"] = layers;
	
	ofs << writer.write(project);
	
	return true;
}

void Stack::weightsInit(double stddev)
{
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		pLayer->weightsInit(stddev);
		pLayer = pLayer->next;
	}
}

bool Stack::loadWeights()
{
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		if (!pLayer->loadWeights(m_name))
		{
			return false;
		}
		pLayer = pLayer->next;
	}
	return true;
}

bool Stack::saveWeights()
{
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		if (!pLayer->saveWeights(m_name))
		{
			return false;
		}
		pLayer = pLayer->next;
	}
	return true;
}

void Stack::train(Rbm::IListener* pListener)
{
	Layer *pLayer = m_pLayers;
	while(pLayer)
	{
		std::cout << m_name << ": " << " Training of layer " << std::to_string(pLayer->id()) << std::endl;
		pLayer->train(trainingData(pLayer), pListener);
		pLayer = pLayer->next;
	}
}

arma::mat& Stack::trainingData()
{
	return m_trainingData;
}

arma::mat Stack::trainingData(Layer* pThatLayer)
{
	arma::mat thisBatch = m_trainingData;
	Layer *pLayer = m_pLayers;
	while (pLayer)
	{
		if (pLayer->id() == pThatLayer->id())
		{
			break;
		}
		thisBatch = pLayer->toHiddenProbs(thisBatch);
		pLayer = pLayer->next;
	}
	return thisBatch;
}

size_t Stack::loadTraining()
{
	uint32_t numTraining = 0;
	uint32_t numVisible = 0;

	std::string filename = m_name + ".training.dat";
	FILE *pFile = fopen(filename.c_str(), "r");

	if (!pFile)
	{
		std::cout << "Could not open " << filename << "!" << std::endl;
		return 0;
	}

	int result = fscanf(pFile, "%d\n", &numTraining);
	if (result < 0)
	{
		return 0;
	}
	result = fscanf(pFile, "%d\n", &numVisible);
	if (result < 0)
	{
		return 0;
	}
	m_trainingData = arma::zeros(numTraining, numVisible);

	uint32_t i, j;
	for (i=0; i < numTraining; i++)
	{
		for (j=0; j < numVisible; j++)
		{
			float v;
			int result = fscanf(pFile, "%f", &v);
			if (result > 0)
			{
				m_trainingData(i, j) = v;
			}
		}
	}
	fclose(pFile);
	return numTraining;
}

size_t Stack::saveTraining()
{
	std::string filename = m_name + ".training.dat";
	FILE *pFile = fopen(filename.c_str(), "w");

	if (!pFile)
	{
		std::cout << "Could not open " << filename << "!" << std::endl;
		return 0;
	}

	fprintf(pFile, "%u\n", (uint32_t)m_trainingData.n_rows);
	fprintf(pFile, "%u\n", (uint32_t)m_trainingData.n_cols);

	uint32_t i, j;

	for (i=0; i < m_trainingData.n_rows; i++)
	{
		for (j=0; j < m_trainingData.n_cols; j++)
		{
			fprintf(pFile, "%3.6f\n", m_trainingData(i, j));
		}
	}

	fclose(pFile);
	
	return m_trainingData.n_rows;
}

size_t Stack::numTraining()
{
	return m_trainingData.n_rows;
}

void Stack::addTraining(const arma::mat &toAdd)
{ 
	m_trainingData.insert_rows(m_trainingData.n_rows, toAdd);
}

void Stack::delTraining(int index)
{
	m_trainingData.shed_row(index);
}

