// Copyright (C) 2003
// Gerhard Neumann (gerhard@igi.tu-graz.ac.at)

//                
// This file is part of RL Toolbox.
// http://www.igi.tugraz.at/ril_toolbox
//
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. The name of the author may not be used to endorse or promote products
//    derived from this software without specific prior written permission.
// 
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
// OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
// IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
// NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
// THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "ril_debug.h"
#include "cpolicygradient.h"
#include "cagent.h"
#include "creinforce.h"

#include <math.h>

CPolicyGradientCalculator::CPolicyGradientCalculator(CAgentController *policy)
{
	this->policy = policy;
}

CGPOMDPGradientCalculator::CGPOMDPGradientCalculator(CRewardFunction *reward, CStochasticPolicy *policy, CAgent *agent, CReinforcementBaseLineCalculator *baseLine, int TStepsPerEpsiode, int Episodes, rlt_real beta) : CPolicyGradientCalculator(policy), CSemiMDPRewardListener(reward)
{
	this->agent = agent;
	this->baseLine = baseLine;

	addParameters(baseLine);
	
	addParameter("GradientEstimationStepsPerEpisode", TStepsPerEpsiode);
	addParameter("GradientEstimationEpisodes", Episodes);
	addParameter("GPOMDPBeta", beta);

	localGradient = new CFeatureList();

	localZTrace = new CFeatureList();

	globalGradient = NULL;

	stochPolicy = policy;
}

CGPOMDPGradientCalculator::~CGPOMDPGradientCalculator()
{
	delete localGradient;
	delete localZTrace;
}

void CGPOMDPGradientCalculator::nextStep(CStateCollection *oldState, CAction *action, rlt_real reward, CStateCollection *newState)
{
	if (globalGradient)
	{
		localZTrace->multFactor(getParameter("GPOMDPBeta"));
		localGradient->clear();

		stochPolicy->getActionProbabilityLnGradient(oldState, action, action->getActionData(),localGradient);
		localZTrace->add(localGradient, 1.0);
		
		CFeatureList::iterator it = localZTrace->begin();
		if (DebugIsEnabled('g'))
		{
			DebugPrint('g', "reward: %f, baseline %f, -> factor %f\n", reward, baseLine->getReinforcementBaseLine((*it)->featureIndex));
			DebugPrint('g', "Z-trace: ");
			localZTrace->saveASCII(DebugGetFileHandle('g'));
		}

		for (;it != localZTrace->end(); it ++)
		{
			globalGradient->update((*it)->featureIndex, (reward - baseLine->getReinforcementBaseLine((*it)->featureIndex)) * (*it)->factor);
		}
	}
}

void CGPOMDPGradientCalculator::newEpisode()
{
	localZTrace->clear();
}

void CGPOMDPGradientCalculator::getGradient(CFeatureList *gradient)
{
	setGlobalGradient(gradient);

	int TSteps = my_round(getParameter("GradientEstimationStepsPerEpisode"));
	int nEpisodes = my_round(getParameter("GradientEstimationEpisodes"));

	agent->startNewEpisode();
	
	bool bListen = agent->isListenerAdded(this);

	if (!bListen)
	{
		agent->addSemiMDPListener(this);
	}

	printf("Calculating PGradient with %d steps and %d Episodes\n", TSteps,nEpisodes);
	
	int oldSteps = 0;
	int gradientSteps = 0;

	oldSteps = agent->getTotalSteps();

	for (int i = 0; i < nEpisodes; i++)
	{
		agent->startNewEpisode();
		agent->doControllerEpisode(1, TSteps);
		printf("Finished %d Episode\n", i);
	}
	gradientSteps = agent->getTotalSteps() - oldSteps;

	if (!bListen)
	{
		agent->removeSemiMDPListener(this);
	}
	assert(gradientSteps > 0);
	gradient->multFactor(1.0 / gradientSteps);

	if (DebugIsEnabled('g'))
	{
		DebugPrint('g', "Calculated GPOMDP Gradient (%d steps)\n", TSteps);
		gradient->saveASCII(DebugGetFileHandle('g'));
		DebugPrint('g', "\n");
	}

	setGlobalGradient(NULL);
}


CFeatureList* CGPOMDPGradientCalculator::getGlobalGradient()
{
	return globalGradient;
}

void CGPOMDPGradientCalculator::setGlobalGradient(CFeatureList *globalGradient)
{
	this->globalGradient = globalGradient;
}

CPolicyGradientUpdater::CPolicyGradientUpdater(CGradientUpdateFunction *updateFunction)
{
	this->updateFunction = updateFunction;
}

void CPolicyGradientUpdater::addRandomParams(rlt_real randSize)
{
	rlt_real *weights = new rlt_real[updateFunction->getNumWeights()];
	updateFunction->getWeights(weights);

	rlt_real normWeights = 0;
	for (int i = 0; i < updateFunction->getNumWeights(); i++)
	{
		normWeights += pow(weights[i], 2);
	}
	normWeights = sqrt(normWeights);

	for (int i = 0; i <updateFunction->getNumWeights(); i ++)
	{
		weights[i] += CDistributions::getNormalDistributionSample(0, normWeights * randSize / 2);
	}
	updateFunction->setWeights(weights);
	delete weights;
}

CConstantPolicyGradientUpdater::CConstantPolicyGradientUpdater(CGradientUpdateFunction *updateFunction, rlt_real learningRate) : CPolicyGradientUpdater(updateFunction)
{
	addParameter("PolicyGradientFactor", learningRate);
}

void CConstantPolicyGradientUpdater::updateWeights(CFeatureList *gradient)
{
	updateFunction->updateGradient(gradient, getParameter("PolicyGradientFactor"));
}


CGSearchPolicyGradientUpdater::CGSearchPolicyGradientUpdater(CGradientUpdateFunction *updateFunction, CPolicyGradientCalculator *gradientCalculator, rlt_real s0, rlt_real epsilon) : CPolicyGradientUpdater(updateFunction)
{
	this->gradientCalculator = gradientCalculator;

	startParameters = new rlt_real[updateFunction->getNumWeights()];
	workParameters = new rlt_real[updateFunction->getNumWeights()];

	addParameters(gradientCalculator, "GSearch");
	addParameter("GSearchStartStepSize", s0);
	addParameter("GSearchEpsilon",epsilon);
	addParameter("GSearchUseLastStepSize", 0.0);

	addParameter("GSearchMinStepSize", s0 / 256);
	addParameter("GSearchMaxStepSize", s0 * 16);

	lastStepSize = s0;
}

CGSearchPolicyGradientUpdater::~CGSearchPolicyGradientUpdater()
{
	delete [] startParameters;
	delete [] workParameters;
}

void CGSearchPolicyGradientUpdater::setWorkingParamters(CFeatureList *gradient, rlt_real stepSize, rlt_real *startParameters, rlt_real *workParameters)
{
	memcpy(workParameters, startParameters, sizeof(rlt_real) * updateFunction->getNumWeights());

	CFeatureList::iterator it = gradient->begin();
	for (; it != gradient->end(); it ++)
	{
		workParameters[(*it)->featureIndex] += stepSize * (*it)->factor;
	}
}

void CGSearchPolicyGradientUpdater::updateWeights(CFeatureList *gradient)
{

	rlt_real s = getParameter("GSearchStartStepSize");

	rlt_real norm = sqrt(gradient->multFeatureList(gradient));

	if (getParameter("GSearchUseLastStepSize") > 0.5)
	{
		s = lastStepSize;
	}
	printf("Beginning GSEARCH with stepSize %f\n", s);

	rlt_real epsilon = getParameter("GSearchEpsilon");

	updateFunction->getWeights(startParameters);
	setWorkingParamters(gradient, s,startParameters, workParameters);

	updateFunction->setWeights(workParameters);
	CFeatureList *newGradient = new CFeatureList();
	gradientCalculator->getGradient(newGradient);

	rlt_real newGradientNorm = sqrt(newGradient->multFeatureList(newGradient));

	rlt_real prod = gradient->multFeatureList(newGradient);// * 1 / newGradientNorm;;
	rlt_real tempProd = prod;
	rlt_real sPlus = 0;
	rlt_real sMinus = 0;
	rlt_real pPlus = 0;
	rlt_real pMinus = 0;

	rlt_real sMin = getParameter("GSearchMinStepSize");
	rlt_real sMax = getParameter("GSearchMaxStepSize");

	printf("gradient * newgradient: %f\n", tempProd);

	if (prod < 0)
	{
		sPlus = s; 

		while(tempProd < - epsilon && s > sMin)
		{
			sPlus = s;
			pPlus = tempProd;
			s = s / 2;

			printf("GSearch StepSize: %f ", s);
			
			setWorkingParamters(gradient, s, startParameters, workParameters);
			updateFunction->setWeights(workParameters);
			newGradient->clear();
			gradientCalculator->getGradient(newGradient);

			newGradientNorm = sqrt(newGradient->multFeatureList(newGradient));
			tempProd = gradient->multFeatureList(newGradient);// * 1 / newGradientNorm;
			
			printf("GSearch StepSize: %f, gradient * newGradient: %f\n", s,tempProd);

		} 
		sMinus = s;
		pMinus = tempProd;
		if (s < sMin)
		{
			s = sMin;
		}
	}
	else
	{
		sMinus = s;
		while(tempProd > epsilon && s < sMax)
		{
			sMinus = s;
			pMinus = tempProd;

			s = 2 * s;

			setWorkingParamters(gradient, s, startParameters, workParameters);
			updateFunction->setWeights(workParameters);
			newGradient->clear();

			gradientCalculator->getGradient(newGradient);

			newGradientNorm = sqrt(newGradient->multFeatureList(newGradient));
			tempProd = gradient->multFeatureList(newGradient);// * 1 / newGradientNorm;

			printf("GSearch StepSize: %f, gradient * newGradient: %f\n", s,tempProd);
		}
		sPlus = s;
		pPlus = tempProd;

		if (s > sMax)
		{
			s = sMax;
		}
	}


	if (pMinus > 0 && pPlus < 0)
	{
		s = (pPlus * sMinus - pMinus * sPlus) / (pPlus - pMinus);
	}
	else
	{
		s = (sPlus + sMinus) / 2;
	}

	printf("GSearch: s: %f, s+ %f, s- %f, p+ %f, p- %f\n",s, sPlus, sMinus, pPlus, pMinus);

	DebugPrint('g',"GSearch: s: %f, s+ %f, s- %f, p+ %f, p- %f\n",s, sPlus, sMinus, pPlus, pMinus);

	setWorkingParamters(gradient, s, startParameters, workParameters);

	if (DebugIsEnabled('g'))
	{
		DebugPrint('g',"GSearch: Calculated StepSize %f\n", s);
		DebugPrint('g', "GSearch: New calculated Parameters\n");
		updateFunction->saveData(DebugGetFileHandle('g'));
	}
	
	lastStepSize = s;


	updateFunction->setWeights(workParameters);

	rlt_real normWeights = 0;
	
	for (int i = 0; i < updateFunction->getNumWeights(); i ++)
	{
		normWeights += workParameters[i] * workParameters[i];
	}
	normWeights = sqrt(normWeights);
	printf("Weights Norm after Update %f\n", normWeights);
}