#include "RooFitter.h"

void RooFitter::FitEMuResults()
{
	gStyle->SetPalette(60); 		// <---<< kBlueRedYellow
	gStyle->SetOptStat(0);
	gStyle->SetOptFit(0);
	const Int_t kSlate = TColor::GetColor(112, 128, 144);
	const Int_t kCrimson = TColor::GetColor(220, 20, 60);
	const Int_t kFirebrick = TColor::GetColor(178, 34, 34);
	const Int_t kPurple = TColor::GetColor(128, 0, 128);
	const Int_t kDarkViolet = TColor::GetColor(148, 0, 211);
	const Int_t kDeepSkyBlue = TColor::GetColor(0, 191, 255);

	for (int i=1; i<=h_FitResult->GetNbinsX(); ++i) {
		double BinContent = h_FitResult->GetBinContent(i);
		double BinError = sqrt(BinContent); 		// <---<< Assuming Poisson uncertainties
		h_FitResult->SetBinError(i, BinError);
	}

	RooRealVar Mass(Branch, Branch, LeftCut, RightCut);
	RooArgSet Variables(Mass);

	RooArgSet NormalizationSet(Mass);

	RooDataHist DataHist("DataHist","Data Histogram",Mass,Import(*h_FitResult));
	RooDataSet* Data = new RooDataSet("Data","Data",RooArgSet(Mass));
	for (int i=0; i<h_FitResult->GetNbinsX(); ++i) {
		Mass.setVal(h_FitResult->GetBinCenter(i+1));
		for (int j=0; j<h_FitResult->GetBinContent(i+1); ++j) {
			Data->add(RooArgSet(Mass));
		}
	}

	/// Functions (Straight Lines) for Residuals
	TF1* F1 = new TF1("F1","1",LeftCut,RightCut);
	F1->SetLineColor(kBlue);
	F1->SetLineStyle(7);

	TF1* F0 = new TF1("F0","0",LeftCut,RightCut);
	F0->SetLineColor(kBlack);

	TF1* Fn1 = new TF1("Fn1","-1",LeftCut,RightCut);
	Fn1->SetLineColor(kBlue);
	Fn1->SetLineStyle(7);

	// Signal Modeling
	#if defined FixMass
		RooRealVar Mean("KPsi_Mean","Mean",PDGMass);
		Mean.setConstant(kTRUE);
	#else
		float MinMean = PDGMass-(PDGMass*0.02);
		float MaxMean = PDGMass+(PDGMass*0.02);
		RooRealVar Mean("KPsi_Mean","Mean",PDGMass,MinMean,MaxMean);
	#endif

	#if defined TripleGaussian
		RooRealVar Sigma1("KPsi_Sigma1","Standard deviation of the first Gaussian P.D.F.", 1.00e-2, 1.00e-3, 1.00e-1);
		RooRealVar Sigma2("KPsi_Sigma2","Standard deviation of the second Gaussian P.D.F.", 1.00e-2, 1.00e-3, 1.00e-1);
		RooRealVar Sigma3("KPsi_Sigma3","Standard deviation of the third Gaussian P.D.F.", 1.00e-2, 1.00e-3, 1.00e-1);

		RooGaussian Gaussian1("KPsi_Gaussian1","First Gaussian P.D.F.",Mass,Mean,Sigma1);
		RooGaussian Gaussian2("KPsi_Gaussian2","Second Gaussian P.D.F.",Mass,Mean,Sigma2);
		RooGaussian Gaussian3("KPsi_Gaussian3","Second Gaussian P.D.F.",Mass,Mean,Sigma3);

		RooRealVar AFC1("KPsi_AFC1","Fraction of first Gaussian P.D.F.", 0.5, 0.0, 1.0); 		// <---<< AFC = Amplitude Fraction Coefficient
		RooRealVar AFC2("KPsi_AFC2","Fraction of second Gaussian P.D.F.", 0.5, 0.0, 1.0);

		RooArgList* PDFList = new RooArgList(Gaussian1,Gaussian2,Gaussian3);
		RooArgList* AFCList = new RooArgList(AFC1,AFC2);
		RooAddPdf Signal("KPsi_Signal","Triple-Gaussian P.D.F.",*PDFList,*AFCList);

	#elif defined BreitWigner
		RooRealVar Gamma("KPsi_Gamma","Width of Breit-Wigner P.D.F.", 1.00e-2, 1.00e-3, 1.00e-1);
		RooBreitWigner Signal("KPsi_Signal","Breit-Wigner P.D.F.",Mass,Mean,Gamma);

	#elif defined Voigt
		RooRealVar Gamma("KPsi_Gamma","Width of Breit-Wigner P.D.F.", 1.00e-2, 1.00e-3, 1.00e-1);
		RooRealVar Sigma("KPsi_Sigma","Standard deviation of Gaussian P.D.F.", 5.00e-3, 1.00e-3, 2.00e-2);
		RooVoigtian Signal("KPsi_Signal","Voigtian P.D.F.",Mass,Mean,Gamma,Sigma);

	#endif

	RooRealVar SignalAmplitude("KPsi_SignalAmplitude","Signal Amplitude", 0, 0, 1.0e6);
	RooExtendPdf ExtSignal("KPsi_ExtSignal","Extended Signal",Signal,SignalAmplitude);

	// Background Modeling
	#if defined Linear

		RooRealVar* Slope = new RooRealVar("Slope","Slope of the First-Order Polynomial Background P.D.F.", 0, -1.0e4, 1.0e4);
		RooRealVar* YIntercept = new RooRealVar("YIntercept","Y-Intercept of the First-Order Polynomial Background P.D.F.", 0, -1.0e4, 1.0e4);
		RooPolynomial Background("PolBackground", "First-Order Polynomial Background P.D.F.",Mass,RooArgList(*YIntercept,*Slope), 1);

	#elif defined Exponential

		RooRealVar cExpo("cExpo","cExpo", 0.0, -1.0e1, 1.0e1);
		RooExponential Background("ExpBackground","Exponential Background P.D.F.",Mass,cExpo);

	#elif defined Polynomial

		RooRealVar A("A","A", 0.1, -1.0e4, 1.0e4);
		RooRealVar B("B","B", -0.1, -1.0e4, 1.0e4);
		RooRealVar C("C","C", 1000, -1.0e4, 1.0e4);
		RooPolynomial Background("PolBackground","Second-Order Polynomial Background P.D.F.",Mass,RooArgList(A, B, C), 2);

	#endif

	RooRealVar BackgroundAmplitude("BackgroundAmplitude","Background Amplitude", 0, 0, 1.0e6);
	RooExtendPdf ExtBackground("ExtBackground","Extended Background",Background,BackgroundAmplitude);

	RooArgList* ExtPDFList = new RooArgList(ExtSignal,ExtBackground);
	RooAddPdf* ModelPDF = new RooAddPdf("ModelPDF","Model P.D.F.",*ExtPDFList);
	ModelPDF->fixCoefNormalization(NormalizationSet);

	RooFitResult* FitResult = ModelPDF->fitTo(DataHist,
																						RooFit::NumCPU(4, kTRUE), 									// <---<< Use 4 CPUs with multi-threading
																						RooFit::Save(), 														// <---<< Save the fit results
																						RooFit::Extended(), 												// <---<< Perform an extended fit
																						RooFit::Minimizer("Minuit", "Minimize"), 		// <---<< Use Minuit for minimization
																						RooFit::Strategy(2), 												// <---<< Set strategy to 2 for more robust fitting
																						RooFit::NormSet(NormalizationSet), 					// <---<< Apply normalization set
																						RooFit::PrintLevel(-1));
	FitResult->Print("v");

	double SignalYield = SignalAmplitude.getVal();
	double SignalYieldError = SignalAmplitude.errorVar()->getVal();

	double BackgroundYield = BackgroundAmplitude.getVal();
	double BackgroundYieldError = BackgroundAmplitude.errorVar()->getVal();

	double TotalYield = SignalYield+BackgroundYield;
	double NumberOfEvents = Data->sumEntries();

	// Fitting PDF's to Multiple Toy Monte Carlo Sets (useful for goodness-of-fit studies)
	#if defined PULLS
		int NMCSets = 1000;

		RooMCStudy* ToyModel = new RooMCStudy(ModelPDF,RooArgSet(Mass),Binned(kFALSE),Silence(),Extended(),FitOptions(Save(kTRUE)),PrintEvalErrors(0),Minos(kTRUE));
		ToyModel->generateAndFit(NMCSets);

		RooPlot* SignalPullFrame = ToyModel->plotPull(SignalAmplitude, -4, 4, 8, kTRUE);
		RooPlot* BackgroundPullFrame = ToyModel->plotPull(BackgroundAmplitude, -4, 4, 8, kTRUE);

		TCanvas* AllPulls = new TCanvas("AllPulls","ToyMC", 800, 800);
		AllPulls->Divide(2,1);
			AllPulls->cd(1);
				SignalPullFrame->Draw();
				SignalPullFrame->SetXTitle("Signal Yield");
			AllPulls->cd(2);
				BackgroundPullFrame->Draw();
				BackgroundPullFrame->SetXTitle("Background Yield");

		FilePath.clear();
		FilePath = OutputDirectoryName + "/AllPullDistributions.png";
		AllPulls->SaveAs(FilePath.c_str());
		FilePath.clear();
		FilePath = OutputDirectoryName + "/AllPullDistributions.C";
		AllPulls->SaveAs(FilePath.c_str());
		FilePath.clear();
	#endif

	// Plotting Data Points
	RooPlot* Frame = Mass.frame(Title(FrameTitle),Bins(NBins),Range(LeftCut,RightCut));
	Frame->SetMinimum(0);
	double MaxHeight = h_FitResult->GetBinContent(h_FitResult->GetMaximumBin());
	Frame->SetMaximum(MaxHeight+(0.75*MaxHeight));
	Data->plotOn(Frame,Name("Data"),XErrorSize(0));

	// Plotting Full Model
	ModelPDF->plotOn(Frame,LineColor(kRed),LineWidth(3),NumCPU(1,kTRUE),Name("Model"));
	//ModelPDF->plotOn(Frame,Components("KPsi_ExtSignal"),LineWidth(3),DrawOption("F"),FillColor(kGreen),LineColor(kGreen),Name("Signal"));
	ModelPDF->plotOn(Frame,Components("KPsi_ExtSignal"),LineWidth(3),LineColor(kGreen),Name("Signal"));
	ModelPDF->plotOn(Frame,Components("ExtBackground"),LineWidth(3),LineColor(kPurple),LineStyle(9),Name("Background"));
	#if defined TripleGaussian
		Signal.plotOn(Frame,Components("KPsi_Gaussian1"),LineWidth(3),LineColor(kFirebrick),Name("Gaussian1"));
		Signal.plotOn(Frame,Components("KPsi_Gaussian2"),LineWidth(3),LineColor(kSlate),Name("Gaussian2"));
		Signal.plotOn(Frame,Components("KPsi_Gaussian3"),LineWidth(3),LineColor(kDarkViolet),Name("Gaussian3"));
	#endif

	/// Creating and Filling Canvas with Plots and Info
	TCanvas* MassProjection = new TCanvas("Mass Fit","Mass Fit", 900, 900);
	MassProjection->SetWindowSize(1920,1080);
	MassProjection->cd();

	/// Extracting Residuals
	RooPlot* ResidualFrame = Mass.frame(Title("Frame to extract residuals"),Bins(NBins),Range(LeftCut,RightCut));
	Data->plotOn(ResidualFrame,XErrorSize(0));
	ModelPDF->plotOn(ResidualFrame);

	RooHist* h_Residuals = ResidualFrame->pullHist();
	RooPlot* MassResidualsFrame = Mass.frame(Title("Residuals"),Range(LeftCut,RightCut));

	MassResidualsFrame->GetXaxis()->SetTitle(""); 		// <---<< No X-Axix Title
	MassResidualsFrame->GetXaxis()->SetLabelSize(0.10);
	MassResidualsFrame->GetYaxis()->SetTitle("Pulls");
	MassResidualsFrame->GetYaxis()->SetTitleSize(0.15);
	MassResidualsFrame->GetYaxis()->SetLabelSize(0.10);
	MassResidualsFrame->GetYaxis()->SetTitleOffset(0.2);
	MassResidualsFrame->GetYaxis()->SetNdivisions(5);
	MassResidualsFrame->addPlotable(h_Residuals,"P");

	#if defined DrawResiduals
		TPad* MainPad = new TPad("MainPad","The pad 80% of the height",0.0,0.2,1.0,1.0,0); 		// <---< Last 0 is the fill color
		MainPad->Draw();
		TPad* ResidualPad = new TPad("ResidualPad","The pad 20% of the height",0.0,0.0,1.0,0.2,0);
		ResidualPad->Draw();
		ResidualPad->cd();
		MassResidualsFrame->Draw();
		F0->Draw("SAME");
		F1->Draw("SAME");
		Fn1->Draw("SAME");
		MainPad->cd();
	#endif

	Frame->GetXaxis()->SetTitleOffset(1.);
	Frame->GetXaxis()->SetTitle(XAxisTitle);
	Frame->GetYaxis()->SetTitle(YAxisTitle);
	Frame->SetMinimum(1e-5);
	Frame->Draw();

	double Offset1 = 0.0;
	double Offset2 = 0.0;
	if (SignalType == "B") {
		Offset1 = 0.0200;
		Offset2 = 0.1575;
	}
	else if (SignalType == "P") {
		Offset1 = 0.050;
		Offset2 = 0.405;
	}
	else if (SignalType == "J") {
		Offset1 = 0.050;
		Offset2 = 0.405;
	}

	TLatex* Latex = new TLatex();
	Latex->SetTextSize(0.03);
	Latex->DrawLatex(LeftCut+Offset1,Frame->GetMaximum()*0.900,"#font[22]{#color[1]{#it{CMS Preliminary}}}");
	Latex->DrawLatex(LeftCut+Offset1,Frame->GetMaximum()*0.850,"#font[22]{#it{#sqrt{s} = 13 TeV, L = 137.19 fb^{-1}}}");
	Latex->DrawLatex(LeftCut+Offset1,Frame->GetMaximum()*0.800,DecayMode);

	Latex->SetTextSize(0.02);
	Latex->DrawLatex(LeftCut+Offset2,Frame->GetMaximum()*0.650,Form("N("+ShortDecayMode+") = %.2f #pm %.2f",SignalYield,SignalYieldError));
	Latex->DrawLatex(LeftCut+Offset2,Frame->GetMaximum()*0.600,Form("N(Background) = %.2f #pm %.2f",BackgroundYield,BackgroundYieldError));
	Latex->DrawLatex(LeftCut+Offset2,Frame->GetMaximum()*0.550,Form("Total number of events returned from the fit = %.2f",TotalYield));
	Latex->DrawLatex(LeftCut+Offset2,Frame->GetMaximum()*0.500,Form("Total number of events = %.0f",NumberOfEvents));

	RooCurve* PlotData = Frame->getCurve("Data");
	RooCurve* PlotModel = Frame->getCurve("Model");
	RooCurve* PlotDG = Frame->getCurve("Signal");
	RooCurve* PlotBG = Frame->getCurve("Background");
	#if defined TripleGaussian
		RooCurve* PlotG1 = Frame->getCurve("Gaussian1");
		RooCurve* PlotG2 = Frame->getCurve("Gaussian2");
		RooCurve* PlotG3 = Frame->getCurve("Gaussian3");
	#endif

	TLegend* MassLeg = new TLegend(0.618, 0.690, 0.830, 0.850);
	MassLeg->SetTextFont(62);
	MassLeg->SetTextSize(0.02);
	MassLeg->SetBorderSize(0);
	MassLeg->SetFillColor(0);
	MassLeg->AddEntry(PlotData,"Data","lep");
	MassLeg->AddEntry(PlotModel,"Total","l");
	//MassLeg->AddEntry(PlotDG,ShortDecayMode,"f");
	MassLeg->AddEntry(PlotDG,ShortDecayMode,"l");
	MassLeg->AddEntry(PlotBG,"Combinatorial Background", "l");
	#if defined TripleGaussian
		MassLeg->AddEntry(PlotG1,"First Gaussian Component","l");
		MassLeg->AddEntry(PlotG2,"Second Gaussian Component","l");
		MassLeg->AddEntry(PlotG3,"Third Gaussian Component","l");
	#endif
	MassLeg->Draw();

	FilePath = OutputDirectoryName+"/Fit.png";
	MassProjection->SaveAs(FilePath.c_str());
	FilePath.clear();

	std::cout << endl;
	TextFile << endl;

	// List of Parameters and Uncertainties
	RooArgSet* Parameters = ModelPDF->getParameters(*Data);
	TIterator* iPar = Parameters->createIterator();
	RooRealVar* Parameter;
	std::vector<std::string> ParameterNames;
	std::vector<double> ParameterValues;
	std::vector<double> ParameterUncertainties;
	while ((Parameter = (RooRealVar*)iPar->Next())) {
		std::cout << "Parameter: " << setw(25) << left << Parameter->GetName() << "\tValue: " << setw(15) << fixed << setprecision(6) << Parameter->getVal();
		std::cout << "\tUncertainty: " << setw(15) << fixed << setprecision(6) << Parameter->getError();
		TextFile << "Parameter: " << setw(25) << left << Parameter->GetName() << "\tValue: " << setw(15) << fixed << setprecision(6) << Parameter->getVal();
		TextFile << "\tUncertainty: " << setw(15) << fixed << setprecision(6) << Parameter->getError();
		if (Parameter->isConstant()) {
			std::cout << " (Fixed)";
			TextFile << " (Fixed)";
		}
		ParameterNames.emplace_back(Parameter->GetName());
		ParameterValues.emplace_back(Parameter->getVal());
		ParameterUncertainties.emplace_back(Parameter->getError());
		std::cout << endl;
		TextFile << endl;
	}
	delete iPar;
	delete Parameters;

	std::cout << endl;
	TextFile << endl;

	// Calculating Chi-Squared over Number of Degrees of Freedom
	RooChi2Var Chi2("Chi2","Chi2",*ModelPDF,DataHist);
	double ChiSquared = Chi2.getVal();
	int NDF = DataHist.numEntries()-FitResult->floatParsFinal().getSize();
	double Chi2NDF = ChiSquared/NDF;
	double FitQuality = FitResult->covQual();
	double FitStatus = FitResult->status();
	double PValue = TMath::Prob(ChiSquared,NDF);
	double SignalSignificance = TMath::NormQuantile(1-PValue);

	std::cout << endl;
	std::cout << "#####################################################################" << endl;
	std::cout << "Chi-Squared over Number of Degrees of Freedom = " << Chi2NDF << endl;
	std::cout << endl;
	std::cout << "Covariance Matrix (Fit) Quality = " << FitQuality << endl;
	std::cout << endl;
	std::cout << "Fit Status = " << FitStatus << endl;
	std::cout << endl;
	std::cout << "P-Value = " << PValue << endl;
	std::cout << endl;
	std::cout << "Significance of Signal = " << SignalSignificance << endl;
	std::cout << "#####################################################################" << endl;
	std::cout << endl;

	TextFile << endl;
	TextFile << "#####################################################################" << endl;
	TextFile << "Chi-Squared over Number of Degrees of Freedom = " << Chi2NDF << endl;
	TextFile << endl;
	TextFile << "Covariance Matrix (Fit) Quality = " << FitQuality << endl;
	TextFile << endl;
	TextFile << "Fit Status = " << FitStatus << endl;
	TextFile << endl;
	TextFile << "P-Value = " << PValue << endl;
	TextFile << endl;
	TextFile << "Significance of Signal = " << SignalSignificance << endl;
	TextFile << "#####################################################################" << endl;
	TextFile << endl;

	FS::path SourceFile = "./FitEMuResults.h";
	FS::path DestinationFile = OutputDirectoryName+"/FitEMuResults.h";

	try {
		FS::copy(SourceFile, DestinationFile, FS::copy_options::overwrite_existing);
		std::cout << "'FitEMuResults.h' copied successfully!" << std::endl;
	}
	catch (FS::filesystem_error& e) {
		std::cerr << "Error copying file: " << e.what() << std::endl;
	}

	SourceFile = "./RooFitter.h";
	DestinationFile = OutputDirectoryName+"/RooFitter.h";

	try {
		FS::copy(SourceFile, DestinationFile, FS::copy_options::overwrite_existing);
		std::cout << "'RooFitter.h' copied successfully!" << std::endl;
	}
	catch (FS::filesystem_error& e) {
		std::cerr << "Error copying file: " << e.what() << std::endl;
	}
}

//----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
// Information:
//
// covQual():
// 		Returns an integer representing the quality of the covariance matrix after the fit.
// 		Typical return values and their interpretations:
// 				3: Covariance matrix is valid (i.e., it is positive definite and has non-zero elements).
// 				2: Covariance matrix is accurate, but not guaranteed to be positive definite.
// 				1: Covariance matrix is not accurate (may have large uncertainties).
// 				0: Covariance matrix is not available or not calculated.
// 		Higher values indicate better quality of the covariance matrix, with 3 being the best.

// status():
// 		Returns an integer representing the status of the fit.
// 		Typical return values and their interpretations:
// 				0: Fit is successful.
// 				1: Hesse matrix is not positive definite.
// 				2: Minimizer reaches maximum number of function calls or iterations without convergence.
// 				3: Fit is unsuccessful due to unknown reason.
// 				Other values may indicate specific issues encountered during the fit.
// 		A return value of 0 typically indicates a successful fit, while non-zero values indicate various problems or issues encountered during the fitting process.
//----------------------------------------------------------------------------------------------------------------------------------------------------------------------------