#include "TCanvas.h"
#include "RooFit.h"
#include "RooArgSet.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooDataSet.h"
#include "RooRealVar.h"

#include "RooChebychev.h"
#include "RooFFTConvPdf.h"
#include "RooNumConvPdf.h"
#include "RooAddPdf.h"
#include "RooGaussian.h"
#include "RooArgusBG.h"

using namespace RooFit;

struct Component
{
        RooAbsPdf &pdf;
        RooRealVar #
};

void MWE_argus()
{
        Double_t Obs_cut[2] = {5500, 6100};
        TString Fitting_Var = "mass";
        TString Fitting_Var_Unit = "MeV";
        TString Fitting_Var_Legend = "#it{m}";
        RooRealVar *Obs = new RooRealVar(Fitting_Var, Fitting_Var_Legend, Obs_cut[0], Obs_cut[1], Fitting_Var_Unit);

        Double_t nentry = 611;

        RooRealVar *mean = new RooRealVar("mean", "mean of Mass", 5800.0, 5795, 5805);
        RooRealVar *sigma = new RooRealVar("sigma", "width of Mass", 8.1, 1, 30);

        vector<Component> PDF_components;

        RooGaussian *sigpdf = new RooGaussian("sigpdf", "gaus", *Obs, *mean, *sigma);
        RooRealVar *nsig = new RooRealVar("nsig", "nsig", 100, 0, nentry);
        PDF_components.push_back({*sigpdf, *nsig});

        RooRealVar *c0 = new RooRealVar("c0", "coefficient #0", -0.004, -1, 1.); /// Coef of Comb. Bkg
        RooExponential *CombBkg = new RooExponential("CombBkg", "Combinatory background component", *Obs, *c0);
        RooRealVar *n_CombBkg = new RooRealVar("n_CombBkg", "n_comb_bkg", 358, 0, nentry);

        PDF_components.push_back({*CombBkg, *n_CombBkg});

        RooRealVar *argus_par = new RooRealVar("argus_par", "argus shape parameter", -55.0);
        RooRealVar *argus_m0 = new RooRealVar("argus_m0", "argus threshold mass", 5800 - 135);
        RooArgusBG *ArgusBkg = new RooArgusBG("ArgusBkg", "Argus PDF", *Obs, *argus_m0, *argus_par);
        RooGaussian *Conv_Resolution = new RooGaussian("Conv_Resolution", "To be Convoluted as the Resolution", *Obs, RooFit::RooConst(0), *sigma);
        RooFFTConvPdf *Part_Rec_Bkg = new RooFFTConvPdf("Part_Rec_Bkg", "Partial Reconstruct Bkg", *Obs, *ArgusBkg, *Conv_Resolution);
        // RooNumConvPdf *Part_Rec_Bkg = new RooNumConvPdf("Part_Rec_Bkg", "Partial Reconstruct Bkg", *Obs, *ArgusBkg, *Conv_Resolution);
        RooRealVar *n_argus = new RooRealVar("n_argus", "yields of argus component", nentry / 2, 0, nentry);
        PDF_components.push_back({*Part_Rec_Bkg, *n_argus});

        // --- Total PDFs ---
        RooArgList shapes, yields;
        for (auto &component : PDF_components)
        {
                shapes.add(component.pdf);
                yields.add(component.num);
        }

        RooAddPdf *model = new RooAddPdf("model", "model", shapes, yields);

        RooDataSet *data = model->generate(*Obs, nentry);

        // RooFitResult *fit_result = model->fitTo(*data, Extended(kTRUE), Save(kTRUE), Minimizer("Minuit2", "Migrad"), NumCPU(1), RooFit::EvalBackend("legacy"));
        RooFitResult *fit_result = model->fitTo(*data, Extended(kTRUE), Save(kTRUE), Minimizer("Minuit2", "Migrad"), NumCPU(1));



        cout << "\033[35m" << endl;
        fit_result->Print("v");
        cout << "\033[0m" << endl;

        TCanvas *can = new TCanvas("can", "", 800, 600);
        RooPlot *frame = Obs->frame(Obs->getMin(), Obs->getMax(), 75);
        data->plotOn(frame);

        model->plotOn(frame, RooFit::Name("Part_Rec_Bkg"), RooFit::Components(*CombBkg, *Part_Rec_Bkg), RooFit::DrawOption("F"), RooFit::FillColor(kYellow), RooFit::LineWidth(0));

        model->plotOn(frame, RooFit::Components(RooArgSet(*CombBkg)), RooFit::DrawOption("F"), RooFit::FillColor(0), RooFit::LineWidth(0)); // clean up the color
        model->plotOn(frame, RooFit::Name("Comb_Bkg"), Components(RooArgSet(*CombBkg)), DrawOption("F"), FillColor(TColor::GetColorTransparent(kCyan, 0.3)), LineWidth(0));

        model->plotOn(frame, RooFit::Name("sig"), Components(RooArgSet(*sigpdf)), LineColor(kRed));

        model->plotOn(frame, RooFit::Name("model"));
        data->plotOn(frame, RooFit::Name("data"));

        frame->Draw();

        system("mkdir -p Figs");
        can->SaveAs("Figs/Mass.pdf");
}