#ifndef DFitControl_h
#define DFitControl_h

#include <string>
#include <vector>
#include <utility>

#include "TObject.h"
#include "TF1.h"
#include "TMatrixDSym.h"

#include "DFitUtilities.h"
#include "DFitFunctor.h"
#include "DFitGroup.h"

using namespace std;

class DFitControl : public TObject
{
	//This class owns all of the input fit functors (when delete is called, it deletes them)
	public:

		class DYieldData
		{
			public:
				double dYield;
				double dYieldUncertainty;
				double dBackground;
				double dBackgroundUncertainty;
				double dNumEvents;
		};

		//Constructors
		DFitControl(double locFitRangeMin, double locFitRangeMax, DFitFunctor* locSignalFitFunctor, string locFitTag = "");
		DFitControl(double locFitRangeMin, double locFitRangeMax, DFitFunctor* locSignalFitFunctor, DFitFunctor* locBackgroundFitFunctor, string locFitTag = "");
		DFitControl(double locFitRangeMin, double locFitRangeMax, vector<DFitFunctor*>& locSignalFitFunctors, vector<DFitFunctor*>& locBackgroundFitFunctors, string locFitTag = "");

		//Post-Fit Constructors
		DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, DFitFunctor* locSignalFitFunctor, string locFitTag = "");
		DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, DFitFunctor* locSignalFitFunctor, DFitFunctor* locBackgroundFitFunctor, string locFitTag = "");
		DFitControl(TF1* locTotalFunc, TMatrixDSym* locCovarianceMatrix, vector<DFitFunctor*>& locSignalFitFunctors, vector<DFitFunctor*>& locBackgroundFitFunctors, string locFitTag = "");

		virtual ~DFitControl(void);

		//Fit Functors
		void Append_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag = true);
		void Replace_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag = true, size_t locFuncIndex = 0); //will delete the old one!
		void Insert_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag = true, size_t locFuncIndex = 0);
		void Remove_FitFunctor(bool locSignalFlag = true, size_t locFuncIndex = 0); //will delete the old one!

		// Evaluate
		double operator()(double* locX, double* locParamArray);

		// Getters
		unsigned int Get_NumParams(void) const;
		unsigned int Get_NumFunctions(void) const;

		unsigned int Get_NumSignalParams(void) const;
		unsigned int Get_NumBackgroundParams(void) const;

		vector<double> Get_InitParams(void) const;

		vector<string> Get_FullParamNames(void) const;
		vector<string> Get_SignalParamNames(void) const;
		vector<string> Get_BackgroundParamNames(void) const;

		void Add_ExcludeDataRange(double locRangeMin, double locRangeMax);

		void Set_CovarianceMatrices(const TMatrixDSym& locCovarianceMatrix);
		void Update_CovarianceMatrices(void);

		template <typename DType> void Remove_FitFunctions(DType* locFitObject);
		template <typename DType> void Add_FitFunctions(DType* locFitObject);

		string Get_ParamStream(void) const;

		void Prepare_Refit(void);

		TF1* Create_TotalFunction(string locFitObjectName);
		void Create_FinalFunctions(string locFitObjectName);

		void Calc_Yield(const TH1* locHist, bool locTrustSignalFlag, bool locYieldOverFitRangeFlag, DYieldData& locYieldData) const;
		void Calc_Yield(const TH1* locHist, bool locTrustSignalFlag, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const;

		double dFitRangeMin;
		double dFitRangeMax;
		bool dBindAboveZeroFlag;
		string dFitTag; //in case fitting a histogram more than once: used for unique TF1 names

		bool dLogLikelihoodFlag;

		TF1* dTotalFunc;
		TMatrixDSym dCovarianceMatrix;

		//add function to hist/graph flags
		bool dAddTotalFuncFlag;

		int dTotalFuncColor;
		unsigned int dNumXPoints;
		int dFitStatus;

		//fit groups
		DFitGroup* dBackgroundFitGroup;
		DFitGroup* dSignalFitGroup;

	protected:
		vector<pair<double, double> > dExcludeDataRanges; //exclude data points within ranges from fit

	private:
		DFitControl(void){}
		void Calc_Yield_TrustSignal(const TH1* locHist, bool locYieldOverRangeFlag, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const;
		void Calc_Yield_TrustBackground(const TH1* locHist, double locRangeMin, double locRangeMax, DYieldData& locYieldData) const;

	ClassDef(DFitControl, 1)
};

template <typename DType> void DFitControl::Remove_FitFunctions(DType* locFitObject)
{
	//WILL NOT delete the functions! Just removes them from the fit object
	if(dAddTotalFuncFlag)
		locFitObject->GetListOfFunctions()->Remove(dTotalFunc);

	dSignalFitGroup->Remove_FitFunctions<DType>(locFitObject);
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->Remove_FitFunctions<DType>(locFitObject);
}

template <typename DType> void DFitControl::Add_FitFunctions(DType* locFitObject)
{
	if(dAddTotalFuncFlag)
		locFitObject->GetListOfFunctions()->Add(dTotalFunc);
	if(Get_NumFunctions() == 1)
		return;

	dSignalFitGroup->Add_FitFunctions<DType>(locFitObject);
	if(dBackgroundFitGroup != NULL)
		dBackgroundFitGroup->Add_FitFunctions<DType>(locFitObject);
}

inline void DFitControl::Append_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag)
{
	if(locSignalFlag)
		dSignalFitGroup->Append_FitFunctor(locFitFunctor);
	else
	{
		if(dBackgroundFitGroup != NULL)
			dBackgroundFitGroup->Append_FitFunctor(locFitFunctor);
		else
			dBackgroundFitGroup = new DFitGroup("Background", locFitFunctor->Get_InitParams(), vector<DFitFunctor*>(1, locFitFunctor), kBlack);
	}
}

inline void DFitControl::Replace_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag, size_t locFuncIndex)
{
	if(locSignalFlag)
		dSignalFitGroup->Replace_FitFunctor(locFitFunctor, locFuncIndex);
	else
	{
		if(dBackgroundFitGroup != NULL)
			dBackgroundFitGroup->Replace_FitFunctor(locFitFunctor, locFuncIndex);
		else
			cout << "WARNING, INVALID REPLACEMENT IN DFITCONTROL::Replace_FitFunctor(). FUNCTOR NOT REPLACED." << endl;
	}
}

inline void DFitControl::Insert_FitFunctor(DFitFunctor* locFitFunctor, bool locSignalFlag, size_t locFuncIndex)
{
	if(locSignalFlag)
		dSignalFitGroup->Insert_FitFunctor(locFitFunctor, locFuncIndex);
	else
	{
		if(dBackgroundFitGroup != NULL)
			dBackgroundFitGroup->Insert_FitFunctor(locFitFunctor, locFuncIndex);
		else
			dBackgroundFitGroup = new DFitGroup("Background", locFitFunctor->Get_InitParams(), vector<DFitFunctor*>(1, locFitFunctor), kBlack);
	}
}

inline void DFitControl::Remove_FitFunctor(bool locSignalFlag, size_t locFuncIndex)
{
	if(locSignalFlag)
		dSignalFitGroup->Remove_FitFunctor(locFuncIndex);
	else
	{
		if(dBackgroundFitGroup != NULL)
			dBackgroundFitGroup->Remove_FitFunctor(locFuncIndex);
		else
			cout << "WARNING, INVALID REPLACEMENT IN DFITCONTROL::Remove_FitFunctor(). FUNCTOR NOT REMOVED." << endl;
	}
}

#endif