#include "DFitControl.h"

ClassImp(DFitControl)

DFitControl::DFitControl(double locFitRangeMin, double locFitRangeMax, DFitFunctor* locSignalFitFunctor, string locFitTag) : 
dFitRangeMin(locFitRangeMin), dFitRangeMax(locFitRangeMax), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000), dBackgroundFitGroup(NULL)
{
	dSignalFitGroup = new DFitGroup("Signal", locSignalFitFunctor->Get_InitParams(), vector<DFitFunctor*>(1, locSignalFitFunctor), 0);
}

DFitControl::DFitControl(double locFitRangeMin, double locFitRangeMax, DFitFunctor* locSignalFitFunctor, DFitFunctor* locBackgroundFitFunctor, string locFitTag) :
dFitRangeMin(locFitRangeMin), dFitRangeMax(locFitRangeMax), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000)
{
	dSignalFitGroup = new DFitGroup("Signal", locSignalFitFunctor->Get_InitParams(), vector<DFitFunctor*>(1, locSignalFitFunctor), 0);
	dBackgroundFitGroup = new DFitGroup("Background", locBackgroundFitFunctor->Get_InitParams(), vector<DFitFunctor*>(1, locBackgroundFitFunctor), 0);
}

DFitControl::DFitControl(double locFitRangeMin, double locFitRangeMax, vector<DFitFunctor*>& locSignalFitFunctors, vector<DFitFunctor*>& locBackgroundFitFunctors, string locFitTag) :
dFitRangeMin(locFitRangeMin), dFitRangeMax(locFitRangeMax), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000)
{
	//setup signal
	vector<double> locInitParams;
	for(size_t loc_i = 0; loc_i < locSignalFitFunctors.size(); ++loc_i)
	{
		vector<double> locFuncParams = locSignalFitFunctors[loc_i]->Get_InitParams();
		locInitParams.insert(locInitParams.end(), locFuncParams.begin(), locFuncParams.end());
	}
	dSignalFitGroup = new DFitGroup("Signal", locInitParams, locSignalFitFunctors, 0);

	//setup background
	locInitParams.clear();
	for(size_t loc_i = 0; loc_i < locBackgroundFitFunctors.size(); ++loc_i)
	{
		vector<double> locFuncParams = locBackgroundFitFunctors[loc_i]->Get_InitParams();
		locInitParams.insert(locInitParams.end(), locFuncParams.begin(), locFuncParams.end());
	}
	dBackgroundFitGroup = locBackgroundFitFunctors.empty() ? NULL : new DFitGroup("Background", locInitParams, locBackgroundFitFunctors, 0);
}

DFitControl::DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, DFitFunctor* locSignalFitFunctor, string locFitTag) : 
dFitRangeMin(0.0), dFitRangeMax(0.0), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000), dBackgroundFitGroup(NULL)
{
	dTotalFunc = locTotalFunc;
	dTotalFunc->GetRange(dFitRangeMin, dFitRangeMax);
	dSignalFitGroup = new DFitGroup("Signal", locTotalFunc, 0, vector<DFitFunctor*>(1, locSignalFitFunctor), kBlack);
	if(locCovarianceMatrix != NULL)
		Set_CovarianceMatrices(*locCovarianceMatrix);
}

DFitControl::DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, DFitFunctor* locSignalFitFunctor, DFitFunctor* locBackgroundFitFunctor, string locFitTag) :
dFitRangeMin(0.0), dFitRangeMax(0.0), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000)
{
	dTotalFunc = locTotalFunc;
	dTotalFunc->GetRange(dFitRangeMin, dFitRangeMax);
	dSignalFitGroup = new DFitGroup("Signal", locTotalFunc, 0, vector<DFitFunctor*>(1, locSignalFitFunctor), kBlack);
	dBackgroundFitGroup = new DFitGroup("Background", locTotalFunc, locSignalFitFunctor->Get_NumParams(), vector<DFitFunctor*>(1, locBackgroundFitFunctor), kBlack);
	if(locCovarianceMatrix != NULL)
		Set_CovarianceMatrices(*locCovarianceMatrix);
}

DFitControl::DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, vector<DFitFunctor*>& locSignalFitFunctors, vector<DFitFunctor*>& locBackgroundFitFunctors, string locFitTag) :
dFitRangeMin(0.0), dFitRangeMax(0.0), dFitTag(locFitTag), 
dLogLikelihoodFlag(false), dAddTotalFuncFlag(true), dTotalFuncColor(kBlack), dNumXPoints(1000)
{
	dTotalFunc = locTotalFunc;
	dTotalFunc->GetRange(dFitRangeMin, dFitRangeMax);
	dSignalFitGroup = new DFitGroup("Signal", locTotalFunc, 0, locSignalFitFunctors, kBlack);
	dBackgroundFitGroup = new DFitGroup("Background", locTotalFunc, dSignalFitGroup->Get_NumParams(), locBackgroundFitFunctors, kBlack);
	if(locCovarianceMatrix != NULL)
		Set_CovarianceMatrices(*locCovarianceMatrix);
}

DFitControl::~DFitControl(void)
{
	delete dTotalFunc;
	delete dSignalFitGroup;
	if(dBackgroundFitGroup != NULL)
		delete dBackgroundFitGroup;
}

// Getters
unsigned int DFitControl::Get_NumParams(void) const
{
	return (Get_NumSignalParams() + Get_NumBackgroundParams());
}

unsigned int DFitControl::Get_NumFunctions(void) const
{
	return (dSignalFitGroup->dFitFunctors.size() + ((dBackgroundFitGroup != NULL) ? dBackgroundFitGroup->dFitFunctors.size() : 0));
}

unsigned int DFitControl::Get_NumSignalParams(void) const
{
	return dSignalFitGroup->Get_NumParams();
}

unsigned int DFitControl::Get_NumBackgroundParams(void) const
{
	return ((dBackgroundFitGroup != NULL) ? dBackgroundFitGroup->Get_NumParams() : 0);
}

vector<double> DFitControl::Get_InitParams(void) const
{
	vector<double> locParams = dSignalFitGroup->Get_InitParams();
	if(dBackgroundFitGroup != NULL)
	{
		vector<double> locBackgroundParams = dBackgroundFitGroup->Get_InitParams();
		locParams.insert(locParams.end(), locBackgroundParams.begin(), locBackgroundParams.end());
	}
	return locParams;
}

vector<string> DFitControl::Get_FullParamNames(void) const
{
	vector<string> locParamNames = Get_SignalParamNames();
	vector<string> locBackgroundParamNames = Get_BackgroundParamNames();
	locParamNames.insert(locParamNames.end(), locBackgroundParamNames.begin(), locBackgroundParamNames.end());
	return locParamNames;
}

vector<string> DFitControl::Get_SignalParamNames(void) const
{
	return dSignalFitGroup->Get_ParamNames();
}

vector<string> DFitControl::Get_BackgroundParamNames(void) const
{
	return ((dBackgroundFitGroup != NULL) ? dBackgroundFitGroup->Get_ParamNames() : vector<string>());
}

void DFitControl::Add_ExcludeDataRange(double locRangeMin, double locRangeMax)
{
	pair<double, double> locPair(locRangeMin, locRangeMax);
	dExcludeDataRanges.push_back(locPair);
	dSignalFitGroup->dExcludeDataRanges.push_back(locPair);
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->dExcludeDataRanges.push_back(locPair);
}

void DFitControl::Set_CovarianceMatrices(const TMatrixDSym& locCovarianceMatrix)
{
	dCovarianceMatrix.ResizeTo(locCovarianceMatrix.GetNcols(), locCovarianceMatrix.GetNrows());
	dCovarianceMatrix = locCovarianceMatrix;
	Update_CovarianceMatrices();
}

void DFitControl::Update_CovarianceMatrices(void)
{
	dSignalFitGroup->Set_CovarianceMatrix(dCovarianceMatrix, 0);
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->Set_CovarianceMatrix(dCovarianceMatrix, Get_NumSignalParams());
}

string DFitControl::Get_ParamStream(void) const
{
	string locStream = dSignalFitGroup->Get_ParamStream();
	if(dBackgroundFitGroup != NULL)
		locStream += dBackgroundFitGroup->Get_ParamStream();
	return locStream;
}

void DFitControl::Prepare_Refit(void)
{
	dSignalFitGroup->Prepare_Refit();
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->Prepare_Refit();
	delete dTotalFunc;
	dTotalFunc = NULL;
}

// Evaluate
double DFitControl::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;
		}
	}

	//signal
	double locValue = (*dSignalFitGroup)(locX, locParamArray);
	unsigned int locParamIndex = Get_NumSignalParams();

	//background
	if(dBackgroundFitGroup != NULL)
		locValue += (*dBackgroundFitGroup)(locX, &(locParamArray[locParamIndex]));

	return locValue;
}

TF1* DFitControl::Create_TotalFunction(string locFitObjectName)
{
	string locFuncName = locFitObjectName;
	if(dFitTag != "")
		locFuncName += string("_") + dFitTag;
	locFuncName += string("_Total");

	//total param array (& errors)
	vector<double> locInitParams = Get_InitParams();
	double* locParamArray = new double[Get_NumParams()];
	for(size_t loc_i = 0; loc_i < locInitParams.size(); ++loc_i)
		locParamArray[loc_i] = locInitParams[loc_i];

	dTotalFunc = new TF1(locFuncName.c_str(), this, dFitRangeMin, dFitRangeMax, Get_NumParams(), "DFitControl");
	dTotalFunc->SetTitle("Total");

	if(Get_NumFunctions() == 1)
	{
		dTotalFunc->SetLineColor(dSignalFitGroup->dFitFunctors[0]->dFuncColor);
		dTotalFunc->SetNpx(dSignalFitGroup->dFitFunctors[0]->dNumXPoints);
	}
	else
	{
		dTotalFunc->SetLineColor(dTotalFuncColor);
		dTotalFunc->SetNpx(dNumXPoints);
	}
	dTotalFunc->SetParameters(locParamArray);

	//FIX PARAMETERS
	{
		//signal
		set<size_t>& locFixedParams = dSignalFitGroup->dFixedParams;
		set<size_t>::iterator locFixedIterator = locFixedParams.begin();
		for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
		{
			size_t locFixParamIndex = *locFixedIterator;
			dTotalFunc->FixParameter(locFixParamIndex, locParamArray[locFixParamIndex]);
		}

		//background
		unsigned int locParamIndex = dSignalFitGroup->Get_NumParams();
		if(dBackgroundFitGroup != NULL)
		{
			locFixedParams = dBackgroundFitGroup->dFixedParams;
			locFixedIterator = locFixedParams.begin();
			for(; locFixedIterator != locFixedParams.end(); ++locFixedIterator)
			{
				size_t locFixParamIndex = locParamIndex + *locFixedIterator;
				dTotalFunc->FixParameter(locFixParamIndex, locParamArray[locFixParamIndex]);
			}
		}
	}

	//SET PARAMETER LIMITS
	{
		//signal
		map<size_t, pair<double, double> >& locParamLimits = dSignalFitGroup->dParamLimits;
		map<size_t, pair<double, double> >::iterator locIterator = locParamLimits.begin();
		for(; locIterator != locParamLimits.end(); ++locIterator)
			dTotalFunc->SetParLimits(locIterator->first, locIterator->second.first, locIterator->second.second);

		//background
		if(dBackgroundFitGroup != NULL)
		{
			unsigned int locParamIndex = dSignalFitGroup->Get_NumParams();
			locParamLimits = dBackgroundFitGroup->dParamLimits;
			for(locIterator = locParamLimits.begin(); locIterator != locParamLimits.end(); ++locIterator)
				dTotalFunc->SetParLimits(locIterator->first + locParamIndex, locIterator->second.first, locIterator->second.second);
		}
	}

	vector<string> locFullParamNames = Get_FullParamNames();
	for(size_t loc_i = 0; loc_i < locFullParamNames.size(); ++loc_i)
	{
		if(locFullParamNames[loc_i] != "")
			dTotalFunc->SetParName(loc_i, locFullParamNames[loc_i].c_str());
	}

	return dTotalFunc;
}

void DFitControl::Create_FinalFunctions(string locFitObjectName)
{
	string locFuncBaseName = locFitObjectName;
	if(dFitTag != "")
		locFitObjectName += string("_") + dFitTag;

	//total param array (& errors)
	double* locParamArray = dTotalFunc->GetParameters();
	double* locParamErrorArray = const_cast<double*>(dTotalFunc->GetParErrors());

	//signal
	unsigned int locNumSignalParams = Get_NumSignalParams();
	dSignalFitGroup->Create_Function(locFuncBaseName, dFitRangeMin, dFitRangeMax, locParamArray, locParamErrorArray);

	//individual background
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->Create_Function(locFuncBaseName, dFitRangeMin, dFitRangeMax, &(locParamArray[locNumSignalParams]), &(locParamErrorArray[locNumSignalParams]));
}

//INTERFACE
void DFitControl::Calc_Yield(const TH1* locHist, bool locTrustSignalFlag, bool locYieldOverFitRangeFlag, DYieldData& locYieldData) const
{
	if(locTrustSignalFlag)
	{
		//Yield from signal integral
			//Range: Fit/Infinite if locYieldOverFitRangeFlag = true/false
			//If infinite is specified but is not possible, fit range is used
		//Background from Total - Yield (Fit Range)
		Calc_Yield_TrustSignal(locHist, locYieldOverFitRangeFlag, dFitRangeMin, dFitRangeMax, locYieldData);
	}
	else
	{
		//Background from background integral (Fit Range)
		//Signal from Total - background (Fit Range)
		//locYieldOverFitRangeFlag ignored
		Calc_Yield_TrustBackground(locHist, dFitRangeMin, dFitRangeMax, locYieldData);
	}
}

//INTERFACE
void DFitControl::Calc_Yield(const TH1* locHist, bool locTrustSignalFlag, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const
{
	if(locTrustSignalFlag)
	{
		//Yield from signal integral (Limited Range)
		//Background from Total - Yield (Limited Range)
		Calc_Yield_TrustSignal(locHist, true, locRangeMin, locRangeMax, locYieldData);
	}
	else
	{
		//Background from background integral (Limited Range)
		//Signal from Total - background (Limited Range)
		Calc_Yield_TrustBackground(locHist, locRangeMin, locRangeMax, locYieldData);
	}
}

void DFitControl::Calc_Yield_TrustSignal(const TH1* locHist, bool locYieldOverRangeFlag, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const
{
	//Yield from signal integral (Limited Range)
	//Background from Total - Yield (Limited Range)

	//Uncertainties:
		//Signal: Described by fit distribution: Error given by fit (integral) error
		//Background: Described by remainder of total distribution: Error given by counting events: Poisson: sqrt(N)

	//Get Bin Width, Num Events
	double locAverageBinWidth = 0.0;
	locYieldData.dNumEvents = DFitUtilities::Calc_EntriesInRange(locHist, locRangeMin, locRangeMax, locAverageBinWidth);

	//Calc Signal Integral
	double locIntegral = 0.0, locIntegralError = 0.0;
	if(locYieldOverRangeFlag)
		dSignalFitGroup->Calc_Integral(locRangeMin, locRangeMax, locIntegral, locIntegralError);
	else //yield from -inf to inf (if possible; else fit range min/max)
		dSignalFitGroup->Calc_Integral(locIntegral, locIntegralError);

	//Calc Yield
	locYieldData.dYield = locIntegral/locAverageBinWidth;
	locYieldData.dYieldUncertainty = locIntegralError/locAverageBinWidth;

	//Calc Background
	locYieldData.dBackground = locYieldData.dNumEvents - locYieldData.dYield;
	locYieldData.dBackgroundUncertainty = sqrt(locYieldData.dBackground);
}

void DFitControl::Calc_Yield_TrustBackground(const TH1* locHist, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const
{
	//Background from background integral (Limited Range)
	//Signal from Total - background (Limited Range)

	//Uncertainties:
		//Background: Described by fit distribution: Error given by fit (integral) error
		//Signal: Described by remainder of total distribution: Error given by counting events: Poisson: sqrt(N)

	//Get Bin Width, Num Events
	double locAverageBinWidth = 0.0;
	locYieldData.dNumEvents = DFitUtilities::Calc_EntriesInRange(locHist, locRangeMin, locRangeMax, locAverageBinWidth);

	//Calc Background Integral
	double locIntegral = 0.0, locIntegralError = 0.0;
	dBackgroundFitGroup->Calc_Integral(locRangeMin, locRangeMax, locIntegral, locIntegralError);

	//Calc Background
	locYieldData.dBackground = locIntegral/locAverageBinWidth;
	locYieldData.dBackgroundUncertainty = locIntegralError/locAverageBinWidth;

	//Calc Yield
	locYieldData.dYield = locYieldData.dNumEvents - locYieldData.dBackground;
	locYieldData.dYieldUncertainty = sqrt(locYieldData.dYield);
}