#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooAddPdf.h"
#include "RooDataSet.h"
#include "RooStats/SPlot.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TTree.h"
#include "TFile.h"

void sweight() {
   RooRealVar mB_Jpsi("mB_Jpsi", "Invariant Mass", 5300, 5900, "MeV/c^{2}");
    RooRealVar mean("mean", "Mean of Gaussian", 5600, 5300, 5700);

    RooRealVar sigma("sigma", "Width of Gaussian", 20, 5, 50);
    RooGaussian signal("signal", "Signal Gaussian", mB_Jpsi, mean, sigma);

    RooRealVar lambda("lambda", "Slope of Exponential", -0.002, -0.01, 0.0);
    RooExponential background("background", "Combinatorial Background", mB_Jpsi, lambda);

    RooRealVar yield_signal("yield_signal", "Signal Yield", 1000, 0, 10000);
    RooRealVar yield_bkg("yield_bkg", "Background Yield", 5000, 0, 20000);

    RooAddPdf model("model", "Signal + Background", RooArgList(signal, background), RooArgList(yield_signal, yield_bkg));

    RooDataSet* data = model.generate(mB_Jpsi, 20000);

    model.fitTo(*data);
    
    cout << "Signal yield (from fit) before splot: " << yield_signal.getVal() << endl;
    cout << "Background yield (from fit) before splot: " << yield_bkg.getVal() << endl;
    
        TCanvas* c = new TCanvas("c", "Fit and sWeights", 800, 600);
    RooPlot* frame = mB_Jpsi.frame();
    data->plotOn(frame);
    model.plotOn(frame);
    model.paramOn(frame);
    model.plotOn(frame, RooFit::Components("background"), RooFit::LineStyle(kDashed), RooFit::LineColor(kRed));
    model.plotOn(frame, RooFit::Components("signal"), RooFit::LineStyle(kDashed), RooFit::LineColor(kGreen));
    frame->Draw();

    c->SaveAs("fit_with_sWeights.png");

 //   RooStats::SPlot* sData = new RooStats::SPlot("sData", "sPlot", *data, &model, RooArgList(yield_signal, yield_bkg));

    // Truncate data to the signal region
    double mass_min = 5450; 
    double mass_max = 5700;

    RooDataSet* data_signal_region = (RooDataSet*)data->reduce(Form("mB_Jpsi > %f && mB_Jpsi < %f", mass_min, mass_max));

RooStats::SPlot* sData = new RooStats::SPlot("sData", "sPlot", *data_signal_region, &model, RooArgList(yield_signal, yield_bkg));

    TFile outputFile("sWeights.root", "RECREATE");
    TTree* tree = new TTree("sWeightsTree", "Tree with sWeights");
    double wt_sig, wt_bkg;
    tree->Branch("mB_Jpsi", &mB_Jpsi, "mB_Jpsi/D");
    tree->Branch("wt_sig", &wt_sig, "wt_sig/D");
    tree->Branch("wt_bkg", &wt_bkg, "wt_bkg/D");

double totalSignalWeight = 0, totalBackgroundWeight = 0;
    for (int i = 0; i < data_signal_region->numEntries(); i++) {
        const RooArgSet* entry = data_signal_region->get(i);
        //mB_Jpsi = entry->getRealValue("mB_Jpsi");
        wt_sig = sData->GetSWeight(i, "yield_signal");
        wt_bkg = sData->GetSWeight(i, "yield_bkg");
        
        totalSignalWeight += sData->GetSWeight(i, "yield_signal");
        totalBackgroundWeight += sData->GetSWeight(i, "yield_bkg");
        tree->Fill();
    }

    tree->Write();
    outputFile.Close();


    
     cout << "Signal yield (from fit): " << yield_signal.getVal() << endl;
    cout << "Background yield (from fit): " << yield_bkg.getVal() << endl;
    
    cout << "Total Signal Weight (sPlot): " << totalSignalWeight << endl;
    cout << "Total Background Weight (sPlot): " << totalBackgroundWeight << endl;
}