#include "DFitGroup.h"

ClassImp(DFitGroup)

DFitGroup::DFitGroup(string locGroupName, vector<double> locInitParams, vector<DFitFunctor*> locFitFunctors, int locFuncColor) : 
DFitFunctor(locGroupName, locInitParams, locFuncColor), dFitFunctors(locFitFunctors), dAddGroupFuncFlag(true)
{
	size_t locParamIndex = 0;
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		if((locFuncColor == 0) && (loc_i == 0))
			dFuncColor = dFitFunctors[loc_i]->dFuncColor;

		//set param names
		vector<string> locFuncParamNames = dFitFunctors[loc_i]->Get_ParamNames();
		for(size_t loc_j = 0; loc_j < locFuncParamNames.size(); ++loc_j)
			dParamNames[locParamIndex + loc_j] = locFuncParamNames[loc_j];

		//fix params
		set<size_t>& locFixedParams = dFitFunctors[loc_i]->dFixedParams;
		set<size_t>::iterator locFixedIterator = locFixedParams.begin();
		for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
			dFixedParams.insert(locParamIndex + *locFixedIterator);

		//set param limits
		map<size_t, pair<double, double> >& locParamLimits = dFitFunctors[loc_i]->dParamLimits;
		map<size_t, pair<double, double> >::iterator locLimitIterator = locParamLimits.begin();
		for(; locLimitIterator != locParamLimits.end(); ++locLimitIterator)
			dParamLimits[locLimitIterator->first + locParamIndex] = locLimitIterator->second;

		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}
}

DFitGroup::DFitGroup(string locGroupName, TF1* locTotalFunc, unsigned int locParamIndex, vector<DFitFunctor*> locFitFunctors, int locFuncColor) : 
DFitFunctor(locGroupName, vector<double>(), locFuncColor), dFitFunctors(locFitFunctors), dAddGroupFuncFlag(true)
{
	double locFitRangeMin, locFitRangeMax;
	locTotalFunc->GetRange(locFitRangeMin, locFitRangeMax);
	string locTotalFuncName = locTotalFunc->GetName();
	string locGroupFuncName = locTotalFuncName.substr(0, locTotalFuncName.size() - 5) + Get_FitType(); //subtract "Total" before adding fit type
	double* locParamArray = &((locTotalFunc->GetParameters())[locParamIndex]);
	double* locParamErrorArray = const_cast<double*>(&((locTotalFunc->GetParErrors())[locParamIndex]));

	DFitFunctor::Create_Function(locGroupFuncName, locFitRangeMin, locFitRangeMax, locParamArray, locParamErrorArray);
}

DFitGroup::~DFitGroup(void)
{
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
		delete dFitFunctors[loc_i];
}

double DFitGroup::operator()(double* locX, double* locParamArray)
{
	//reject points entirely if desired
	for(size_t loc_i = 0; loc_i < dExcludeDataRanges.size(); ++loc_i)
	{
		if((locX[0] > dExcludeDataRanges[loc_i].first) && (locX[0] < dExcludeDataRanges[loc_i].second))
		{
			TF1::RejectPoint();
			return 0.0;
		}
	}

	//add func values
	double locValue = 0.0;
	unsigned int locParamIndex = 0;
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		locValue += (*dFitFunctors[loc_i])(locX, &(locParamArray[locParamIndex]));
		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}

	return locValue;
}

void DFitGroup::Set_CovarianceMatrix(const TMatrixDSym& locFullCovarianceMatrix, unsigned int locParamIndex)
{
	DFitFunctor::Set_CovarianceMatrix(locFullCovarianceMatrix, locParamIndex);
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		dFitFunctors[loc_i]->Set_CovarianceMatrix(locFullCovarianceMatrix, locParamIndex);
		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}
}

void DFitGroup::Create_Function(string locFuncBaseName, double locFitRangeMin, double locFitRangeMax, double* locParamArray, double* locParamErrorArray)
{
	string locGroupFuncName = locFuncBaseName + string("_") + Get_FitType();
	DFitFunctor::Create_Function(locGroupFuncName, locFitRangeMin, locFitRangeMax, locParamArray, locParamErrorArray);

	unsigned int locParamIndex = 0;
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		ostringstream locFuncName;
		locFuncName << locFuncBaseName << "_" << Get_FitType() << "_" << loc_i;

		dFitFunctors[loc_i]->Create_Function(locFuncName.str(), locFitRangeMin, locFitRangeMax, &(locParamArray[locParamIndex]), &(locParamErrorArray[locParamIndex]));
		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}
}

void DFitGroup::Prepare_Refit(void)
{
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		for(size_t loc_j = 0; loc_j < dFitFunctors[loc_i]->Get_NumParams(); ++loc_j)
			dFitFunctors[loc_i]->dInitParams[loc_j] = dFitFunctors[loc_i]->dFunc->GetParameter(loc_j);
		delete dFitFunctors[loc_i]->dFunc;
		dFitFunctors[loc_i]->dFunc = NULL;
	}
	for(size_t loc_i = 0; loc_i < Get_NumParams(); ++loc_i)
		dInitParams[loc_i] = dFunc->GetParameter(loc_i);
	delete dFunc;
	dFunc = NULL;
}

void DFitGroup::Calc_Integral(double& locIntegral, double& locIntegralError) const
{
	locIntegral = 0.0;
	locIntegralError = 0.0;
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		double locFunctionIntegral = 0.0, locFunctionIntegralError = 0.0;
		dFitFunctors[loc_i]->Calc_Integral(locFunctionIntegral, locFunctionIntegralError);
		locIntegral += locFunctionIntegral;
		locIntegralError += locFunctionIntegralError*locFunctionIntegralError;
	}
	locIntegralError = sqrt(locIntegralError);
}

void DFitGroup::Calc_Integral(double locRangeMin, double locRangeMax, double& locIntegral, double& locIntegralError) const
{
	locIntegral = 0.0;
	locIntegralError = 0.0;
	for(size_t loc_i = 0; loc_i < dFitFunctors.size(); ++loc_i)
	{
		double locFunctionIntegral = 0.0, locFunctionIntegralError = 0.0;
		dFitFunctors[loc_i]->Calc_Integral(locRangeMin, locRangeMax, locFunctionIntegral, locFunctionIntegralError);
		locIntegral += locFunctionIntegral;
		locIntegralError += locFunctionIntegralError*locFunctionIntegralError;
	}
	locIntegralError = sqrt(locIntegralError);
}

void DFitGroup::Remove_FitFunctor(size_t locFuncIndex)
{
	if(locFuncIndex >= dFitFunctors.size())
	{
		cout << "WARNING, INVALID INDEX IN DFitGroup::Remove_FitFunctor(). FUNCTOR NOT REMOVED." << endl;
		return;
	}

	size_t locPreRemoveParamIndex = 0;
	for(size_t loc_i = 0; loc_i < locFuncIndex; ++loc_i)
		locPreRemoveParamIndex += dFitFunctors[loc_i]->Get_NumParams();

	unsigned int locNumParamsRemoved = dFitFunctors[locFuncIndex]->Get_NumParams();

	//remove param names
	vector<string>::iterator locNameIterator = dParamNames.begin() + locPreRemoveParamIndex;
	for(size_t loc_i = 0; loc_i < dFitFunctors[locFuncIndex]->Get_NumParams(); ++loc_i)
		locNameIterator = dParamNames.erase(locNameIterator);

	//Remove any previously-fixed params
	set<size_t>& locFixedParams = dFitFunctors[locFuncIndex]->dFixedParams;
	set<size_t>::iterator locFixedIterator = locFixedParams.begin();
	for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
		dFixedParams.erase(locPreRemoveParamIndex + *locFixedIterator);

	//Adjust indices of any fixed params from later functors
	size_t locParamIndex = locPreRemoveParamIndex;
	for(size_t loc_i = locFuncIndex + 1; loc_i < dFitFunctors.size(); ++loc_i)
	{
		locFixedParams = dFitFunctors[loc_i]->dFixedParams;
		locFixedIterator = locFixedParams.begin();
		for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
		{
			dFixedParams.erase(locParamIndex + locNumParamsRemoved + *locFixedIterator);
			dFixedParams.insert(locParamIndex + *locFixedIterator);
		}
		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}

	//Remove any previously-set param limits
	map<size_t, pair<double, double> >& locParamLimits = dFitFunctors[locFuncIndex]->dParamLimits;
	map<size_t, pair<double, double> >::iterator locLimitIterator = locParamLimits.begin();
	for(; locLimitIterator != locParamLimits.end(); ++locLimitIterator)
		dParamLimits.erase(locPreRemoveParamIndex + locLimitIterator->first);

	//Adjust indices of any param limits from later functors
	locParamIndex = locPreRemoveParamIndex;
	for(size_t loc_i = locFuncIndex + 1; loc_i < dFitFunctors.size(); ++loc_i)
	{
		locParamLimits = dFitFunctors[loc_i]->dParamLimits;
		locLimitIterator = locParamLimits.begin();
		for(; locLimitIterator != locParamLimits.end(); ++locLimitIterator)
		{
			dParamLimits.erase(locParamIndex + locNumParamsRemoved + locLimitIterator->first);
			dParamLimits[locParamIndex + locLimitIterator->first] = locLimitIterator->second;
		}
		locParamIndex += dFitFunctors[loc_i]->Get_NumParams();
	}

	delete dFitFunctors[locFuncIndex];
	dFitFunctors.erase(dFitFunctors.begin() + locFuncIndex);
}

void DFitGroup::Insert_FitFunctor(DFitFunctor* locFitFunctor, size_t locFuncIndex)
{
	if(locFuncIndex > dFitFunctors.size())
	{
		cout << "WARNING, INVALID INDEX IN DFitGroup::Insert_FitFunctor(). FUNCTOR NOT INSERTED." << endl;
		return;
	}

	size_t locPreInsertParamIndex = 0;
	for(size_t loc_i = 0; loc_i < locFuncIndex; ++loc_i)
		locPreInsertParamIndex += dFitFunctors[loc_i]->Get_NumParams();

	unsigned int locNumParamsInserted = locFitFunctor->Get_NumParams();

	//Insert the param names
	vector<string> locFuncParamNames = locFitFunctor->Get_ParamNames();
	dParamNames.insert(dParamNames.begin() + locPreInsertParamIndex, locFuncParamNames.begin(), locFuncParamNames.end());

	//Adjust indices of any fixed params from later functors
	size_t locParamIndex = Get_NumParams();
	for(int loc_i = int(dFitFunctors.size() - 1); loc_i >= int(locFuncIndex); --loc_i)
	{
		set<size_t>& locFixedParams = dFitFunctors[loc_i]->dFixedParams;
		set<size_t>::reverse_iterator locReverseFixedIterator = locFixedParams.rbegin();
		for(; locReverseFixedIterator != locFixedParams.rend(); ++locReverseFixedIterator)
		{
			dFixedParams.erase(locParamIndex + *locReverseFixedIterator);
			dFixedParams.insert(locParamIndex + locNumParamsInserted + *locReverseFixedIterator);
		}
		locParamIndex -= dFitFunctors[loc_i]->Get_NumParams();
	}

	//Fix params
	set<size_t>& locFixedParams = locFitFunctor->dFixedParams;
	set<size_t>::iterator locFixedIterator = locFixedParams.begin();
	for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
		dFixedParams.insert(locParamIndex + *locFixedIterator);

	//Adjust indices of any param limits from later functors
	locParamIndex = Get_NumParams();
	for(int loc_i = int(dFitFunctors.size() - 1); loc_i >= int(locFuncIndex); --loc_i)
	{
		map<size_t, pair<double, double> >& locParamLimits = dFitFunctors[loc_i]->dParamLimits;
		map<size_t, pair<double, double> >::reverse_iterator locReverseLimitIterator = locParamLimits.rbegin();
		for(; locReverseLimitIterator != locParamLimits.rend(); ++locReverseLimitIterator)
		{
			dParamLimits.erase(locParamIndex + locReverseLimitIterator->first);
			dParamLimits[locParamIndex + locNumParamsInserted + locReverseLimitIterator->first] = locReverseLimitIterator->second;
		}
		locParamIndex -= dFitFunctors[loc_i]->Get_NumParams();
	}

	//set param limits
	map<size_t, pair<double, double> >& locParamLimits = locFitFunctor->dParamLimits;
	map<size_t, pair<double, double> >::iterator locLimitIterator = locParamLimits.begin();
	for(; locLimitIterator != locParamLimits.end(); ++locLimitIterator)
		dParamLimits[locLimitIterator->first + locParamIndex] = locLimitIterator->second;

	//Insert the functor
	dFitFunctors.insert(dFitFunctors.begin() + locFuncIndex, locFitFunctor);
}

void DFitGroup::Append_FitFunctor(DFitFunctor* locFitFunctor)
{
	size_t locParamIndex = Get_NumParams();

	//set param names
	vector<string> locFuncParamNames = locFitFunctor->Get_ParamNames();
	dParamNames.insert(dParamNames.begin() + locParamIndex, locFuncParamNames.begin(), locFuncParamNames.end());

	//fix params
	set<size_t>& locFixedParams = locFitFunctor->dFixedParams;
	set<size_t>::iterator locFixedIterator = locFixedParams.begin();
	for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
		dFixedParams.insert(locParamIndex + *locFixedIterator);

	//set param limits
	map<size_t, pair<double, double> >& locParamLimits = locFitFunctor->dParamLimits;
	map<size_t, pair<double, double> >::iterator locLimitIterator = locParamLimits.begin();
	for(; locLimitIterator != locParamLimits.end(); ++locLimitIterator)
		dParamLimits[locLimitIterator->first + locParamIndex] = locLimitIterator->second;
}