#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 30 10:35:41 2023

@author: yannickburkard
"""


import ROOT



#------------------------------------------------------------------------------------------------------------

#Higgs transverse mass variable
ROOT.gInterpreter.Declare('''
using Vec_t = const ROOT::RVec<float>&;
using namespace ROOT::Math;
float ComputeMT4(ROOT::Math::PtEtaPhiMVector& p1, ROOT::Math::PtEtaPhiMVector& p2, ROOT::Math::PtEtaPhiMVector& p3, ROOT::Math::PtEtaPhiMVector& p4) {
    return sqrt(pow((p1.Pt()+p2.Pt()+p3.Pt()+p4.Pt()),2)-pow((p1+p2+p3+p4).Pt(),2));
}
''')

#smearing function
ROOT.gInterpreter.Declare('''
#include <cmath>
#include <Math/Vector4D.h>
using namespace ROOT::Math;
TRandom2 *randvar = new TRandom2();
PtEtaPhiMVector smear_MET(ROOT::Math::PtEtaPhiMVector& MET, Double32_t N) {
    Double32_t mag = 1.*(randvar->Rndm());
    Double32_t angle = 2*M_PI*(randvar->Rndm());
    auto dpx = mag*cos(angle);
    auto dpy = mag*sin(angle);
    const PxPyPzEVector MET_P4(MET.Px()+dpx,MET.Py()+dpy,MET.Pz(),sqrt(pow((MET.Px()+dpx),2) + pow((MET.Py()+dpy),2) + pow(MET.Pz(),2)));
    const PtEtaPhiMVector MET_return(MET_P4.Pt(), MET_P4.Eta(), MET_P4.Phi(), 0);
    return MET_return;
}
''')


#------------------------------------------------------------------------------------------------------------

dff = ROOT.RDataFrame("Tree", "example_file.root")

N=30.0
dff = dff.Define('N',str(N))
dff = dff.Redefine('MET','smear_MET(MET, N)')
###dff = dff.Redefine('MET','MET') #even redefining MET as itself produces an object of type ROOT::Math::LorentzVector<ROOT::Math::PtEtaPhiM4D<double> >.
dff.Describe().Print()

dff = dff.Define('mHT','ComputeMT4(L,MET,Jet1,Jet2)')
countFilter = dff.Filter('fabs((L+MET).Eta()) < 1.2').Count()
print(countFilter.GetValue())

hist_mHT = dff.Histo1D(("hist_mHT", "hist", 10, 0, 100),"mHT")
mTH_np = dff.AsNumpy(["mHT"])["mHT"]