#include "TROOT.h"
#include "TFile.h"
#include "TTree.h"
#include "TH1F.h"
#include "TH2D.h"
#include "TLorentzVector.h"
#include <iostream>
#include <vector>

bool inSelection(const TLorentzVector& v) {
    return v.E() > 110
            && v.Pz() < 90
            && v.Pt() < 80 && v.Pt() > 65;
}

void exc4() {
    using namespace std;

    Double_t px, py, pz, E;
    TFile* f = new TFile("signal.root");
    TTree* tree;
    f->GetObject("stree", tree);
    auto px_branch = tree->GetBranch("px"); px_branch->SetAddress(&px);
    auto py_branch = tree->GetBranch("py"); py_branch->SetAddress(&py);
    auto pz_branch = tree->GetBranch("pz"); pz_branch->SetAddress(&pz);
    auto e_branch = tree->GetBranch("E"); e_branch->SetAddress(&E);

    Double_t bpx, bpy, bpz, bE;
    TFile* f2 = new TFile("background.root");
    TTree* tree2;
    f2->GetObject("btree", tree2);
    auto bpx_branch = tree2->GetBranch("px"); bpx_branch->SetAddress(&bpx);
    auto bpy_branch = tree2->GetBranch("py"); bpy_branch->SetAddress(&bpy);
    auto bpz_branch = tree2->GetBranch("pz"); bpz_branch->SetAddress(&bpz);
    auto be_branch = tree2->GetBranch("E"); be_branch->SetAddress(&bE);


    TH1D* signal = new TH1D("spx", "Signal E", 100, 0, 0);
    TH1D* background = new TH1D("bpx", "Background py", 100, 0, 0);

    TH2D* signal2 = new TH2D("s2d", "Signal 2D", 100, 0, 0, 100, 0, 0);
    TH2D* background2 = new TH2D("b2d", "Background 2D", 100, 0, 0, 100, 0, 0);

    vector<TLorentzVector> events, selection;
    assert(tree->GetEntries() == tree2->GetEntries());
    int backgrounds = 0, signals = 0;
    int N = tree->GetEntries();
    for (int i = 0; i < N; ++i) {
        px_branch->GetEntry(i);
        py_branch->GetEntry(i);
        pz_branch->GetEntry(i);
        e_branch->GetEntry(i);

        bpx_branch->GetEntry(i);
        bpy_branch->GetEntry(i);
        bpz_branch->GetEntry(i);
        be_branch->GetEntry(i);

        TLorentzVector vsignal;
        TLorentzVector vbackground;

        vsignal.SetPxPyPzE(px, py, pz, E);
        vbackground.SetPxPyPzE(bpx, bpy, bpz, bE);

        signal->Fill(vsignal.E());
        background->Fill(vbackground.E());

        signal2->Fill(vsignal.Pt(), vsignal.Eta());
        background2->Fill(vbackground.Pt(), vbackground.Pt());

        events.push_back(vsignal);
        events.push_back(vbackground);
        if (inSelection(vsignal)) {
            selection.push_back(vsignal);
        } else {
            signals++;
        }
        if (inSelection(vbackground)) {
            backgrounds++;
            selection.push_back(vbackground);
        }
    }
    cout << "Selection: " << selection.size() << endl
        << "False positive: " << backgrounds << endl
        << "False negative: " << signals << endl;

    signal2->SetLineColor(kRed);
    signal2->SetFillColor(kRed);
    signal2->SetMarkerColor(kRed);
    signal2->Draw();
    background2->Draw("Same");
}

int main() {
    exc4();
}