/*
  ==============================================================================

  This is an automatically generated GUI class created by the Introjucer!

  Be careful when adding custom code to these files, as only the code within
  the "//[xyz]" and "//[/xyz]" sections will be retained when the file is loaded
  and re-saved.

  Created with Introjucer version: 3.1.0

  ------------------------------------------------------------------------------

  The Introjucer is part of the JUCE library - "Jules' Utility Class Extensions"
  Copyright 2004-13 by Raw Material Software Ltd.

  ==============================================================================
*/


#include "RbmComponent.hpp"



//==============================================================================
RbmComponent::RbmComponent (DeepStack &stack, const std::string &name, size_t id, size_t numVisibleX, size_t numVisibleY, size_t numHidden, size_t numContext)
	: Layer(name, id, numVisibleX, numVisibleY, numHidden, numContext)
	, m_stack(stack)
	, m_currWeightIndexToDraw(0)
	, DrawVisibleTrain(nullptr)
	, DrawVisibleReconst(nullptr)
	, DrawWeights(nullptr)
	, DrawHidden(nullptr)
	, DrawContextReconst(nullptr)
	, m_sizeX(430)
	, m_sizeY(190)
{
	addAndMakeVisible (DrawVisibleTrain = new DrawComponent (numVisibleX, numVisibleY));
	DrawVisibleTrain->setListener(this);

	addAndMakeVisible (DrawHidden = new DrawComponent (numHidden, 1));
	DrawHidden->setListener(this);

	addAndMakeVisible (DrawContextTrain = new DrawComponent (numContext, 1));
	DrawContextTrain->setListener(this);

	addAndMakeVisible (DrawContextReconst = new DrawComponent (numContext, 1));
	DrawContextReconst->setListener(this);

	addAndMakeVisible (DrawVisibleReconst = new DrawComponent (numVisibleX, numVisibleY));
	addAndMakeVisible (DrawWeights = new DrawComponent (numVisibleX, numVisibleY, 0.5, 0.5));

	addAndMakeVisible (m_toggleEnable = new ToggleButton ("Enable layer"));
	m_toggleEnable->setTooltip (TRANS("Enable layer for up pass"));
	m_toggleEnable->setButtonText (TRANS("Enable"));
	m_toggleEnable->addListener (this);
	m_toggleEnable->setToggleState(true, NotificationType::dontSendNotification);
	
    addAndMakeVisible (m_buttonCopyH2C = new TextButton ("Copy hidden state to context"));
    m_buttonCopyH2C->setTooltip (TRANS("Copy hidden state to context input"));
    m_buttonCopyH2C->setButtonText (TRANS("H2C"));
    m_buttonCopyH2C->addListener (this);
	
	redrawWeights();

	setSize (m_sizeX, m_sizeY);
	resized();
	setBounds (16, (m_sizeY + 10)*id+16, m_sizeX, m_sizeY);
}

RbmComponent::~RbmComponent()
{
	DrawVisibleTrain = nullptr;
	DrawVisibleReconst = nullptr;
	DrawWeights = nullptr;
	DrawHidden = nullptr;
	DrawContextTrain = nullptr;
	DrawContextReconst = nullptr;
	m_buttonCopyH2C = nullptr;
}

//==============================================================================
void RbmComponent::paint (Graphics& g)
{
	g.fillAll (Colours::white);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Sigma"),
				32, 306, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Lambda"),
				120, 306, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Num. epochs"),
				32, 354, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Alpha"),
				120, 354, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Sparsity"),
				208, 306, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Sigma decay"),
				296, 306, 96, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Weight decay"),
				296, 354, 96, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Momentum"),
				208, 354, 80, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Sparsity Alpha"),
				408, 306, 96, 14,
				Justification::centred, true);

	g.setColour (Colours::black);
	g.setFont (Font (15.00f, Font::plain));
	g.drawText (TRANS("Weight init"),
				408, 354, 96, 14,
				Justification::centred, true);

}

void RbmComponent::resized()
{
	DrawVisibleTrain->setBoundsRelative(0, 0, 100./m_sizeX, 100./m_sizeY);
	DrawVisibleReconst->setBoundsRelative(110./m_sizeX, 0, 100./m_sizeX, 100./m_sizeY);
	DrawWeights->setBoundsRelative (220./m_sizeX, 0, 100./m_sizeX, 100./m_sizeY);
	DrawContextTrain->setBoundsRelative (0, 110./m_sizeY, 430./m_sizeX, 20./m_sizeY);
	DrawContextReconst->setBoundsRelative (0, 140./m_sizeY, 430./m_sizeX, 20./m_sizeY);
	DrawHidden->setBoundsRelative (0, 170./m_sizeY, 430./m_sizeX, 20./m_sizeY);
	m_toggleEnable->setBoundsRelative (330./m_sizeX, 0, 80./m_sizeX, 24./m_sizeY);
	m_buttonCopyH2C->setBoundsRelative (330./m_sizeX, 30./m_sizeY, 80./m_sizeX, 24./m_sizeY);
}

void RbmComponent::mouseMove (const MouseEvent& e)
{
    //[UserCode_mouseMove] -- Add your code here...
    //[/UserCode_mouseMove]
}

void RbmComponent::mouseEnter (const MouseEvent& e)
{
    //[UserCode_mouseEnter] -- Add your code here...
    //[/UserCode_mouseEnter]
}

void RbmComponent::mouseExit (const MouseEvent& e)
{
    //[UserCode_mouseExit] -- Add your code here...
    //[/UserCode_mouseExit]
}

void RbmComponent::mouseDown (const MouseEvent& e)
{
    //[UserCode_mouseDown] -- Add your code here...
    //[/UserCode_mouseDown]
}

void RbmComponent::mouseDrag (const MouseEvent& e)
{
    //[UserCode_mouseDrag] -- Add your code here...
//	e.eventComponent->setCentrePosition(e.getPosition().x, e.getPosition().y);
	std::cout << "Mouse = " << e.x << "," << e.y << std::endl;
    //[/UserCode_mouseDrag]
}

void RbmComponent::mouseUp (const MouseEvent& e)
{
    //[UserCode_mouseUp] -- Add your code here...
    //[/UserCode_mouseUp]
}

void RbmComponent::mouseDoubleClick (const MouseEvent& e)
{
    //[UserCode_mouseDoubleClick] -- Add your code here...
    //[/UserCode_mouseDoubleClick]
}

void RbmComponent::mouseWheelMove (const MouseEvent& e, const MouseWheelDetails& wheel)
{
    //[UserCode_mouseWheelMove] -- Add your code here...
    //[/UserCode_mouseWheelMove]
}

void RbmComponent::onDraw(DrawComponent &obj)
{
	if (&obj == DrawHidden)
	{
		m_stack.downPass(id(), obj.getData());
	}
	if (&obj == DrawVisibleTrain)
	{
		m_stack.upDownPass(id(), getTraining());
	}
	if (&obj == DrawContextTrain)
	{
		m_stack.upDownPass(id(), getTraining());
	}
}

void RbmComponent::buttonClicked(Button* buttonThatWasClicked)
{
	if (buttonThatWasClicked == m_toggleEnable)
	{
		bool state = buttonThatWasClicked->getToggleState();
		setEnable(state);

		if (state)
		{
			if (prev)
			{
				RbmComponent *pComp = static_cast<RbmComponent*> (prev);
				pComp->m_toggleEnable->setToggleState(state, NotificationType::sendNotification);
			}
		}
		else
		{
			if (next)
			{
				RbmComponent *pComp = static_cast<RbmComponent*> (next);
				pComp->m_toggleEnable->setToggleState(state, NotificationType::sendNotification);
			}
		}
		redrawReconstruction();
	} 
	else if (buttonThatWasClicked == m_buttonCopyH2C)
	{
		DrawContextTrain->getData() = DrawHidden->getData();
		m_stack.upDownPass(id(), getTraining());
	}
}

void RbmComponent::redrawReconstruction()
{
	RbmComponent *pComp = static_cast<RbmComponent*> (root());
	m_stack.upDownPass(pComp->id(), pComp->getTraining());
}

arma::mat RbmComponent::getTraining() const
{
	return arma::join_rows(DrawVisibleTrain->getData(), DrawContextTrain->getData());
}

arma::mat RbmComponent::getReconst() const
{
	return arma::join_rows(DrawVisibleReconst->getData(), DrawContextReconst->getData());
}

void RbmComponent::trainRedraw(const arma::mat& vc)
{
	DrawVisibleTrain->getData() = vc_to_v(vc);
	DrawContextTrain->getData() = vc_to_c(vc);
	DrawVisibleTrain->DrawData();
	DrawContextTrain->DrawData();
}

void RbmComponent::reconstRedraw(const arma::mat& vc)
{
	DrawVisibleReconst->getData() = vc_to_v(vc);
	DrawContextReconst->getData() = vc_to_c(vc);
	DrawVisibleReconst->DrawData();
	DrawContextReconst->DrawData();
}

void RbmComponent::onUpPass(const arma::mat& v)
{
	DrawVisibleTrain->getData() = vc_to_v(v);
	DrawVisibleTrain->DrawData();
	DrawContextTrain->getData() = vc_to_c(v);
	DrawContextTrain->DrawData();
	trainRedraw(v);
}

void RbmComponent::onDownPass(const arma::mat& h)
{
	DrawHidden->getData() = h;
	DrawHidden->DrawData();

	arma::mat r = prob(h_to_v(h));
	reconstRedraw(r);
}

arma::mat RbmComponent::getConvolutedWeight(arma::mat const &w)
{
//	if (next)
//	{
//		arma::mat wc = w * next->w().t(); // * 1.0/sqrt((double)w.cols());
//		return static_cast<RbmComponent*>(next)->getConvolutedWeight(wc);
//	}
//	else
	{
		return w;
	}
	
}
void RbmComponent::redrawWeights()
{
	DrawWeights->getData() = getConvolutedWeight(whv().col(m_currWeightIndexToDraw));
	DrawWeights->DrawData();
}

void RbmComponent::redrawWeights(size_t index)
{
	m_currWeightIndexToDraw = index;
	redrawWeights();
}

void RbmComponent::setTrainingData(arma::mat const& batch)
{
	RbmComponent *pComp = static_cast<RbmComponent*> (root());
	m_stack.upDownPass(id(), batch);
}
//[/MiscUserCode]


//==============================================================================