#include <iostream>
#include <sstream>
#include <string>
#include <fstream>
#include <map>
#include <set>

#include "PlotDrawer/DCanvasInstructions.h"
#include "PlotDrawer/DPadInstructions.h"
#include "PlotDrawer/DPlotDrawer.h"

#include "PlotFitter/DPlotFitter.h"
#include "PlotFitter/DFitControl.h"
#include "PlotFitter/DFitUtilities.h"
#include "PlotFitter/DFitFunctors_Miscellaneous.h"
#include "PlotFitter/DFitFunctors_Gaussians.h"
#include "PlotFitter/DFitFunctors_BreitWigners.h"

#include "PlotData/DStatUtilities.h"

#include "particleType.h"

double Calc_InitGaussHeight(TH1* locMassHist, double locGaussMean, double locGaussWidth, double locNumSigmas);

DFitFunctor* Build_PolyBackground_Linear(TH1* locMassHist, double locFitRangeMin, double locFitRangeMax);
DFitFunctor* Build_PolyBackground_Quad(TH1* locMassHist, double locGaussMean, double locGaussWidth, double locNumSigmas, double locFitRangeMin, double locFitRangeMax);

/*
.x $ROOT_ANALYSIS_HOME/Load.C
.L scripts_analysis/Fit_Lambda1520.C+
Draw_Dalitz("final/kmkp.root");
*/
void Draw_Dalitz(string locInputFileName)
{
	DPlotDrawer locPlotDrawer;
	DPlotFitter locPlotFitter;
	locPlotFitter.dQuietFitFlag = true;

	TFile* locInputFile = new TFile(locInputFileName.c_str(), "READ");

	DCanvasInstructions* locCanvasInstructions = new DCanvasInstructions("Heya");
	locCanvasInstructions->Set_TargetCanvasAspectRatio(16.0/9.0);
	locCanvasInstructions->dTargetPadAspectRatio = 16.0/9.0;
	locCanvasInstructions->dMinPadWidth = 400;
	locCanvasInstructions->dStyle->SetOptFit(102);
	locCanvasInstructions->dStyle->SetOptStat(10);
	locCanvasInstructions->dStyle->SetPadLeftMargin(0.10);
	locCanvasInstructions->dPadLargeRightMargin = 0.1;
	locCanvasInstructions->dPadInstructions[0]->dLogZFlag = true;

	TH2* locDalitzHist = (TH2*)locInputFile->Get("Hist_Dalitz_PostKinFitCut_KinFit/Dalitz");
	locDalitzHist->GetXaxis()->SetTitle("K^{#plus}K^{#minus} Invariant Mass^{2} (GeV/c^{2})^{2}");
	locDalitzHist->GetYaxis()->SetTitle("pK^{#minus} Invariant Mass^{2} (GeV/c^{2})^{2}");
	locPlotDrawer.Draw_Object(locDalitzHist, locCanvasInstructions);

	locCanvasInstructions->dPadInstructions[0]->dLogXFlag = true;
	TH2* locRhoPhiHist = (TH2*)locInputFile->Get("Hist_InvariantMass_Phi_KinFit/InvariantMassVsConfidenceLevel_LogX");
	locRhoPhiHist->GetXaxis()->SetRangeUser(1.0E-40, 0.0);
	locPlotDrawer.Draw_Object(locRhoPhiHist, locCanvasInstructions);

//	TH1* locPhiHist = locRhoPhiHist->ProjectionY("PhiMass", locRhoPhiHist->GetXaxis()->FindBin(5.73304E-7), locRhoPhiHist->GetNbinsX());
	TH1* locPhiHist = locRhoPhiHist->ProjectionY("PhiMass", locRhoPhiHist->GetXaxis()->FindBin(1.0001E-4), locRhoPhiHist->GetNbinsX());
	TH1* locPhiHist_All = locRhoPhiHist->ProjectionY("PhiMass_All", 1, locRhoPhiHist->GetNbinsX());
	locCanvasInstructions->dPadInstructions[0]->dLogXFlag = false;
	locCanvasInstructions->dPadInstructions[0]->dLogZFlag = false;
	locCanvasInstructions->dPadInstructions[0]->dLegendNames.push_back("All");
	locCanvasInstructions->dPadInstructions[0]->dLegendNames.push_back("KinFit Cut: 10^{-4}");
	TObjArray* locPhiArray = new TObjArray();
	locPhiArray->AddLast(locPhiHist_All);
	locPhiArray->AddLast(locPhiHist);
	locPlotDrawer.Draw_Superimpose(locPhiArray, locCanvasInstructions);
}

/*
.x ../Load.C
.L ~/Code/Work/gluex/ROOT_Scripts/strange_v2/Fit_Lambda1520.C+
Draw_Phis();
*/
void Draw_Phis(void)
{
	DPlotDrawer locPlotDrawer;
	DPlotFitter locPlotFitter;
	locPlotFitter.dQuietFitFlag = true;

	string locInputFileName_40 = "~/Scratch/gluex/strange_v2/2kp_40cut.root";
	TFile* locInputFile_40 = new TFile(locInputFileName_40.c_str(), "READ");
	string locInputFileName_4 = "~/Scratch/gluex/strange_v2/2kp_4cut.root";
	TFile* locInputFile_4 = new TFile(locInputFileName_4.c_str(), "READ");

	DCanvasInstructions* locCanvasInstructions = new DCanvasInstructions("Heya");
	locCanvasInstructions->Set_TargetCanvasAspectRatio(16.0/9.0);
	locCanvasInstructions->dTargetPadAspectRatio = 16.0/9.0;
	locCanvasInstructions->dMinPadWidth = 400;
	locCanvasInstructions->dStyle->SetOptFit(102);
	locCanvasInstructions->dStyle->SetOptStat(10);

	TH1* locPhiHist_40 = (TH1*)locInputFile_40->Get("Hist_InvariantMass_Phi_PostKinFitCut_KinFit/InvariantMass");
	locPlotDrawer.Draw_Object(locPhiHist_40, locCanvasInstructions);
/*
	TH1* locPhiHist_4 = (TH1*)locInputFile_4->Get("Hist_InvariantMass_Phi_PostKinFitCut_KinFit/InvariantMass");
	locCanvasInstructions->dPadInstructions[0]->dLegendNames.push_back("KinFit Cut: 10^{-40}");
	locCanvasInstructions->dPadInstructions[0]->dLegendNames.push_back("KinFit Cut: 10^{-4}");
	TObjArray* locPhiArray = new TObjArray();
	locPhiArray->AddLast(locPhiHist_40);
	locPhiArray->AddLast(locPhiHist_4);
	locPlotDrawer.Draw_Superimpose(locPhiArray, locCanvasInstructions);
*/
}

/*
.x $ROOT_ANALYSIS_HOME/Load.C
.L scripts_analysis/Fit_Lambda1520.C+
Calc_Yields("final/kmkp.root");
*/
void Calc_Yields(string locInputFileName)
{
	DPlotDrawer locPlotDrawer;
	DPlotFitter locPlotFitter;
	locPlotFitter.dQuietFitFlag = true;

	TFile* locInputFile = new TFile(locInputFileName.c_str(), "READ");

	DCanvasInstructions* locCanvasInstructions = new DCanvasInstructions("Heya");
	locCanvasInstructions->dStyle->SetOptFit(102);
	locCanvasInstructions->dStyle->SetOptStat(10);
	locCanvasInstructions->dStyle->SetPadLeftMargin(0.13);
	locCanvasInstructions->dStyle->SetTitleOffset(1.1, "Y");

	TH1* locMassHist = (TH1*)locInputFile->Get("Hist_InvariantMass_YStar_PostKinFitCut_KinFit/InvariantMass");
	locMassHist->GetXaxis()->SetRangeUser(1.4, 2.0);

	locPlotDrawer.Draw_Object(locMassHist, locCanvasInstructions);
return;
	vector<double> locInitParams;
	double locFitRangeMin = 1.465, locFitRangeMax = 1.64;

	//Background: LINEAR
	double locSlope = 0.0, locIntercept = 0.0;
	DFitUtilities::Guess_Params_Linear(locMassHist, locFitRangeMin, locFitRangeMax, locIntercept, locSlope, true);
	DFitFunctor* locPolynomial = Build_PolyBackground_Linear(locMassHist, locFitRangeMin, locFitRangeMax);
/*
	//Background: QUADRATIC
	locInitParams.resize(3);
	//locInitParams[0] = locIntercept;  locInitParams[1] = locSlope;  locInitParams[2] = 0.0;
	locInitParams[0] = 1.23E4;  locInitParams[1] = -2.009E4;  locInitParams[2] = 8044.0;
	locPolynomial = new DFitFunctor_Polynomial(locInitParams, kMagenta);
*/
	//Signal
	//Lambda(1520) -> p, K-
	//p J^P = 1/2+, K- J^P = 0-, Lambda(1520) J^P = 3/2-
	//So by J's: L = 1 or 2
	//If L = 0, J^P = 1/2-
	//If L = 1, J^P = 1/2+, 3/2+
	//If L = 2, J^P = 3/2-, 5/2- //THIS IS IT
	unsigned int locL = 2;
	unsigned int locNumBodies = 3; //# final state particles

	double locTargetMass = ParticleMass(Proton);
	double locAvgBeamE = 8.8; //Beam E is 8.4 -> 9.05: Choose ~avg (8.8)
	//W = sqrt(s) = sqrt(Total_E^2 - Total_P^2)
	//W = sqrt((locAvgBeamE + locTargetMass)*(locAvgBeamE + locTargetMass) - locAvgBeamE*locAvgBeamE)
	//W = sqrt(2.0*locAvgBeamE*locTargetMass + locTargetMass*locTargetMass)
	double locW = sqrt(2.0*locAvgBeamE*locTargetMass + locTargetMass*locTargetMass);

	/*
	//FOR Phase space weight!
	vector<double> locProductMasses(4, 0.0);
	locProductMasses[0] = ParticleMass(KPlus);  locProductMasses[1] = 1.51954; //Lambda(1520)
	locProductMasses[2] = ParticleMass(Proton);  locProductMasses[3] = ParticleMass(KMinus);
*/
	vector<double> locProductMasses(2, 0.0);
	locProductMasses[0] = ParticleMass(Proton);  locProductMasses[1] = ParticleMass(KMinus);

	locInitParams.resize(3); //height, mean, FWHM
	double locSignalBackground = locSlope*1.520 + locIntercept;
//	double locSignalBackground = 230.0;
	double locSignalHeight = locMassHist->GetBinContent(locMassHist->FindBin(1.520)) - locSignalBackground;
//	locSignalHeight *= 1000.0;
	locInitParams[0] = locSignalHeight;  locInitParams[1] = 1.520;  locInitParams[2] = 0.020;
//	double* locParamArray = new double[3];
//	locParamArray[0] = 13674.9;  locParamArray[1] = 1.520;  locParamArray[2] = 0.020;

	DFitFunctor* locBreitWigner = new DFitFunctor_RelBreitWignerGammaDepends(locInitParams, locL, locProductMasses, kBlue);
	//DFitFunctor* locBreitWigner = new DFitFunctor_PhaseSpaceRelBreitWignerGammaDepends(locInitParams, locL, locProductMasses, locNumBodies, locW, locTargetMass, kBlue);
//	TF1* locFunc = new TF1("BW", locBreitWigner, locFitRangeMin, locFitRangeMax, 3, "DFitFunctor");
//	locFunc->SetParameters(locParamArray);

	DFitControl* locFitControl = new DFitControl(locFitRangeMin, locFitRangeMax, locBreitWigner, locPolynomial);
	locFitControl->dLogLikelihoodFlag = false;

	//Do fit
	int locFitStatus = locPlotFitter.Fit_Hist(locMassHist, locFitControl);

	locPlotDrawer.Draw_Object(locMassHist, locCanvasInstructions);
//	locFunc->Draw("SAME");

	DFitControl::DYieldData locYieldData;
	locFitControl->Calc_Yield(locMassHist, false, true, locYieldData);
	cout << "Yield: " << locYieldData.dYield << " +/- " << 100.0*locYieldData.dYieldUncertainty/locYieldData.dYield << "%" << endl;
}

/*
.x Load.C
.L ~/Code/Work/gluex/ROOT_Scripts/strange_v2/Fit_Lambda1520.C+
Fit_Lambda1520("~/Scratch/gluex/strange_v2/2kp.root");
*/
void Fit_Lambda1520(string locInputFileName)
{
	DPlotDrawer locPlotDrawer;
	DPlotFitter locPlotFitter;
	locPlotFitter.dQuietFitFlag = true;

	TFile* locInputFile = new TFile(locInputFileName.c_str(), "READ");

	DCanvasInstructions* locCanvasInstructions = new DCanvasInstructions("Heya");
	locCanvasInstructions->Set_TargetCanvasAspectRatio(16.0/9.0);
	locCanvasInstructions->dTargetPadAspectRatio = 16.0/9.0;
	locCanvasInstructions->dMinPadWidth = 400;
	locCanvasInstructions->dStyle->SetOptFit(102);

	TH2* loc2DHist = (TH2*)locInputFile->Get("BestCombo/YStarMassVsT");
	TH1* locMassHist = loc2DHist->ProjectionY("MassProjTBin", loc2DHist->GetXaxis()->FindBin(-0.9999999), loc2DHist->GetXaxis()->FindBin(-0.100001));

	int locRebinAmount = 12;
	//Sideband phi
	loc2DHist = (TH2*)locInputFile->Get("BestCombo/ProdPlanePhiVsT_Sideband_PARA");
	loc2DHist->Rebin2D(1, locRebinAmount);
	TH1* locProdPlanePhiHist_Sideband_PARA = loc2DHist->ProjectionY("ProdPlanePhi_Sideband_PARA", loc2DHist->GetXaxis()->FindBin(-0.9999999), loc2DHist->GetXaxis()->FindBin(-0.100001));
	locProdPlanePhiHist_Sideband_PARA->Sumw2();
	locPlotDrawer.Draw_Object(locProdPlanePhiHist_Sideband_PARA, locCanvasInstructions);

	loc2DHist = (TH2*)locInputFile->Get("BestCombo/ProdPlanePhiVsT_Sideband_PERP");
	loc2DHist->Rebin2D(1, locRebinAmount);
	TH1* locProdPlanePhiHist_Sideband_PERP = loc2DHist->ProjectionY("ProdPlanePhi_Sideband_PERP", loc2DHist->GetXaxis()->FindBin(-0.9999999), loc2DHist->GetXaxis()->FindBin(-0.100001));
	locProdPlanePhiHist_Sideband_PERP->Sumw2();
	locPlotDrawer.Draw_Object(locProdPlanePhiHist_Sideband_PERP, locCanvasInstructions);

	//Signal phi
	loc2DHist = (TH2*)locInputFile->Get("BestCombo/ProdPlanePhiVsT_Signal_PARA");
	loc2DHist->Rebin2D(1, locRebinAmount);
	TH1* locProdPlanePhiHist_Signal_PARA = loc2DHist->ProjectionY("ProdPlanePhi_Signal_PARA", loc2DHist->GetXaxis()->FindBin(-0.9999999), loc2DHist->GetXaxis()->FindBin(-0.100001));
	locProdPlanePhiHist_Signal_PARA->Sumw2();

	loc2DHist = (TH2*)locInputFile->Get("BestCombo/ProdPlanePhiVsT_Signal_PERP");
	loc2DHist->Rebin2D(1, locRebinAmount);
	TH1* locProdPlanePhiHist_Signal_PERP = loc2DHist->ProjectionY("ProdPlanePhi_Signal_PERP", loc2DHist->GetXaxis()->FindBin(-0.9999999), loc2DHist->GetXaxis()->FindBin(-0.100001));
	locProdPlanePhiHist_Signal_PERP->Sumw2();

	locMassHist->Rebin(2);
	vector<double> locInitParams;
	double locFitRangeMin = 1.465, locFitRangeMax = 1.64;

	//Background: LINEAR
	double locSlope = 0.0, locIntercept = 0.0;
	DFitUtilities::Guess_Params_Linear(locMassHist, locFitRangeMin, locFitRangeMax, locIntercept, locSlope, true);
	DFitFunctor* locPolynomial = Build_PolyBackground_Linear(locMassHist, locFitRangeMin, locFitRangeMax);
/*
	//Background: QUADRATIC
	locInitParams.resize(3);
	locInitParams[0] = locIntercept;  locInitParams[1] = locSlope;  locInitParams[2] = 0.0;
	locPolynomial = new DFitFunctor_Polynomial(locInitParams, kMagenta);
*/
	//Signal
	//Lambda(1520) -> p, K-
	//p J^P = 1/2+, K- J^P = 0-, Lambda(1520) J^P = 3/2-
	//So by J's: L = 1 or 2
	//If L = 0, J^P = 1/2-
	//If L = 1, J^P = 1/2+, 3/2+
	//If L = 2, J^P = 3/2-, 5/2- //THIS IS IT
	unsigned int locL = 2;
	unsigned int locNumBodies = 3; //# final state particles

	double locTargetMass = ParticleMass(Proton);
	double locAvgBeamE = 8.8; //Beam E is 8.4 -> 9.05: Choose ~avg (8.8)
	//W = sqrt(s) = sqrt(Total_E^2 - Total_P^2)
	//W = sqrt((locAvgBeamE + locTargetMass)*(locAvgBeamE + locTargetMass) - locAvgBeamE*locAvgBeamE)
	//W = sqrt(2.0*locAvgBeamE*locTargetMass + locTargetMass*locTargetMass)
	double locW = sqrt(2.0*locAvgBeamE*locTargetMass + locTargetMass*locTargetMass);

	vector<double> locProductMasses(4, 0.0);
	locProductMasses[0] = ParticleMass(KPlus);  locProductMasses[1] = 1.51954; //Lambda(1520)
	locProductMasses[2] = ParticleMass(Proton);  locProductMasses[3] = ParticleMass(KMinus);

	locInitParams.resize(3); //height, mean, FWHM
	double locSignalBackground = locSlope*1.520 + locIntercept;
//	double locSignalBackground = 230.0;
	double locSignalHeight = locMassHist->GetBinContent(locMassHist->FindBin(1.520)) - locSignalBackground;
//	locSignalHeight *= 1000.0;
	locInitParams[0] = locSignalHeight;  locInitParams[1] = 1.520;  locInitParams[2] = 0.020;
//	double* locParamArray = new double[3];
//	locParamArray[0] = 13674.9;  locParamArray[1] = 1.520;  locParamArray[2] = 0.020;

	DFitFunctor* locBreitWigner = new DFitFunctor_PhaseSpaceRelBreitWignerGammaDepends(locInitParams, locL, locProductMasses, locNumBodies, locW, locTargetMass, kBlue);
//	TF1* locFunc = new TF1("BW", locBreitWigner, locFitRangeMin, locFitRangeMax, 3, "DFitFunctor");
//	locFunc->SetParameters(locParamArray);

	DFitControl* locFitControl = new DFitControl(locFitRangeMin, locFitRangeMax, locBreitWigner, locPolynomial);
	locFitControl->dLogLikelihoodFlag = false;

	//Do fit
	int locFitStatus = locPlotFitter.Fit_Hist(locMassHist, locFitControl);

	locPlotDrawer.Draw_Object(locMassHist, locCanvasInstructions);
//	locFunc->Draw("SAME");

	double locSidebandIntegral = locMassHist->Integral(locMassHist->GetXaxis()->FindBin(1.600001), locMassHist->GetXaxis()->FindBin(1.6799999), "width");
	double locPeakBackgroundIntegral = locPolynomial->dFunc->Integral(1.48, 1.58);
	double locSidebandScaleFactor = -1.0*locPeakBackgroundIntegral/locSidebandIntegral;
cout << "side integral, peak back interal, factor = " << locSidebandIntegral << ", " << locPeakBackgroundIntegral << ", " << locSidebandScaleFactor << endl;
	//Subtract sidebands
	locProdPlanePhiHist_Signal_PARA->Add(locProdPlanePhiHist_Sideband_PARA, locSidebandScaleFactor);
	locProdPlanePhiHist_Signal_PERP->Add(locProdPlanePhiHist_Sideband_PERP, locSidebandScaleFactor);
	locPlotDrawer.Draw_Object(locProdPlanePhiHist_Signal_PARA, locCanvasInstructions);
	locPlotDrawer.Draw_Object(locProdPlanePhiHist_Signal_PERP, locCanvasInstructions);

	//Ratio: (PARA - PERP)/(PARA + PERP)
	TH1* locRatioHist = (TH1*)locProdPlanePhiHist_Signal_PARA->Clone("Ratio");
	locRatioHist->Add(locProdPlanePhiHist_Signal_PERP, -1.0); //numerator
	TH1* locDenominatorHist = (TH1*)locProdPlanePhiHist_Signal_PARA->Clone("Ratio");
	locDenominatorHist->Add(locProdPlanePhiHist_Signal_PERP, 1.0); //denominator
	locRatioHist->Divide(locRatioHist, locDenominatorHist);

	//Fit ratio
	locInitParams.resize(2);
	double locAmplitude = locRatioHist->GetMaximum();
	locInitParams[0] = locAmplitude;  locInitParams[1] = 0.0;
	DFitFunctor_Cosine* locCosine = new DFitFunctor_Cosine(locInitParams, true, kRed);

	locFitControl = new DFitControl(-180.0, 180.0, locCosine, "Cosine");
	locFitControl->dLogLikelihoodFlag = true;
	locFitStatus = locPlotFitter.Fit_Hist(locRatioHist, locFitControl);

	locPlotDrawer.Draw_Object(locRatioHist, locCanvasInstructions);
}

DFitFunctor* Build_PolyBackground_Linear(TH1* locMassHist, double locFitRangeMin, double locFitRangeMax)
{
	double locSlope = 0.0, locIntercept = 0.0;
	DFitUtilities::Guess_Params_Linear(locMassHist, locFitRangeMin, locFitRangeMax, locIntercept, locSlope, true);

	vector<double> locInitParams;
	locInitParams.resize(2);
	locInitParams[0] = locIntercept;  locInitParams[1] = locSlope;
	return (new DFitFunctor_LinearAboveZero(locInitParams, locFitRangeMin, locFitRangeMax, kMagenta));
}

DFitFunctor* Build_PolyBackground_Quad(TH1* locMassHist, double locGaussMean, double locGaussWidth, double locNumSigmas, double locFitRangeMin, double locFitRangeMax)
{
	double locSlope = 0.0, locIntercept = 0.0;
	DFitUtilities::Guess_Params_Linear(locMassHist, locFitRangeMin, locFitRangeMax, locIntercept, locSlope, true);
/*
//max of quad just to side of gauss
//y = mx + b
//Delta-y = mx_big - mx_small
	double locFlatParam, locLinearParam, locQuadParam;
	double locCriticalPointX = 0.5*(locFitRangeMax + locGaussMean + locNumSigmas*locGaussWidth); //avg
	double locCriticalPointY = locMassHist->GetBinContent(locMassHist->FindBin(locCriticalPointX));
	double locCurve_VerticalChange = locSlope*(locFitRangeMax - locFitRangeMin);
	double locCurve_HorizontalRangeOfVerticalChange = locFitRangeMax - locFitRangeMin;
*/

//max of quad add edge of hist
//y = mx + b
//Delta-y = mx_big - mx_small
	double locFlatParam, locLinearParam, locQuadParam;
	double locCriticalPointX = locMassHist->GetBinCenter(locMassHist->GetNbinsX());
	double locCriticalPointY = locMassHist->GetBinContent(locMassHist->FindBin(locCriticalPointX));
	double locCurve_VerticalChange = locSlope*(locFitRangeMax - locFitRangeMin);
	double locCurve_HorizontalRangeOfVerticalChange = locFitRangeMax - locFitRangeMin;

	DFitUtilities::Calc_ParabolaParams(locCurve_VerticalChange, locCurve_HorizontalRangeOfVerticalChange, locCriticalPointX, locCriticalPointY, locFlatParam, locLinearParam, locQuadParam);

	vector<double> locInitParams;
	locInitParams.resize(3);
	locInitParams[0] = locFlatParam;  locInitParams[1] = locLinearParam;  locInitParams[2] = locQuadParam;
	return (new DFitFunctor_Polynomial(locInitParams, kMagenta));
}