import root_pandas as rp
import numpy as np
from uproot3_methods import TLorentzVectorArray as lorentzarr
from uproot3_methods import TLorentzVector as lorentz
from tqdm import tqdm
import sys
from shutil import copyfile
import os
from ROOT import TFile, TTree, TList

#fname = sys.argv[1]
#tname = sys.argv[2]

def makeNewKstar(fname,tname,MC):
    d = rp.read_root(fname+'.root', key=tname)

    # load DTF momenta in arrays of 4-vectors
    fourmom = {}
    for pi in ['mu_plus', 'mu_minus', 'pi_zero_resolved', 'K_plus']:
        fourmom[pi] = lorentzarr.from_cartesian(d[pi+'_PX_DTF'], d[pi+'_PY_DTF'], d[pi+'_PZ_DTF'], d[pi+'_PE_DTF'])
        
    # calculate the rescaling (k)
    Bplus_M_PDG = 5279.34
    kvalues = []
    for i in tqdm(range(len(d)), desc='calculating pi0 rescaling'):
        ks = np.linspace(0.2, 2, 1000)
        ppi0 = fourmom['pi_zero_resolved'][i]
        pi0_rescaled = lorentzarr.from_xyzm(ppi0.x*ks, ppi0.y*ks, ppi0.z*ks, fourmom['pi_zero_resolved'][i].mass)
        Bmass_rescaled = (pi0_rescaled + fourmom['K_plus'][i] + fourmom['mu_plus'][i] + fourmom['mu_minus'][i]).mass
        bestk = ks[np.argmin(abs(Bmass_rescaled-Bplus_M_PDG))]
        kvalues.append(bestk)
        
    # now apply the rescaling
    pi0_p3 = fourmom['pi_zero_resolved'].p3 * kvalues
    fourmom['pi0_rescaled'] = lorentzarr.from_xyzm(pi0_p3.x, pi0_p3.y, pi0_p3.z, fourmom['pi_zero_resolved'].mass)

    d['pi_zero_resolved_PX_scaled'] = fourmom['pi0_rescaled'].x
    d['pi_zero_resolved_PY_scaled'] = fourmom['pi0_rescaled'].y
    d['pi_zero_resolved_PZ_scaled'] = fourmom['pi0_rescaled'].z
    d['pi_zero_resolved_PT_scaled'] = fourmom['pi0_rescaled'].pt
    d['pi_zero_resolved_PE_scaled'] = fourmom['pi0_rescaled'].E
    d['K_star_plus_M_scaled'] = (fourmom['pi0_rescaled'] + fourmom['K_plus']).mass

    if(MC):
        truefourmom = {}
        for pi in ['K_star_plus']:
            truefourmom[pi] = lorentzarr.from_cartesian(d[pi+'_TRUEP_X'], d[pi+'_TRUEP_Y'], d[pi+'_TRUEP_Z'], d[pi+'_TRUEP_E'])
            
        d['K_star_plus_M_TRUE'] = truefourmom['K_star_plus'].mass

    d.to_root(fname+'_pi0rescaled.root', key=tname, store_index=False)
    return



def cutFile(fname, tname):
    d = rp.read_root(fname+'_pi0rescaled.root', key=tname,where='K_star_plus_M_scaled>791.66 & K_star_plus_M_scaled<991.66') #It takes & not &&!!!!!!!!
    d.to_root(fname+'_pi0rescaledCut.root', key=tname, store_index=False)
    copyfile(fname+'_pi0rescaledCut.root',fname+'.root')
    os.remove(fname+'_pi0rescaledCut.root')
    
def mergeFiles(fname1, fname2, outname, tname):
    tfile1 = TFile(fname1, 'read')
    tfile2 = TFile(fname2, 'read')
    ttree1 = tfile1.Get(tname)
    ttree2 = tfile2.Get(tname)
    treeList = TList()
    treeList.append(ttree1)
    treeList.append(ttree2)
    outFile = TFile(outname,'recreate')
    outFile.cd()
    outTree = TTree.MergeTrees(treeList)
    outFile.Write()
    outFile.Close()
    return


years = [2011,2012,2015,2016,2017,2018]
treeName = "DecayTree"
mainPath = "/home/lhcb/kopecna/B2KstarMuMu_clean/Data/Tuples"
#MC
for pol in ["down","up"]:
    for yr in years:
        if (yr==2015): continue
        fileName = mainPath+"/MC/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
        makeNewKstar(fileName,treeName,True)
        cutFile(fileName,treeName)

#Ref
for pol in ["down","up"]:
    for yr in years:
        if (yr>2016): continue
        fileName = mainPath+"/RefMC/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
        makeNewKstar(fileName,treeName,True)
        cutFile(fileName,treeName)

#PHSP
for pol in ["down","up"]:
    for yr in years:
        fileName = mainPath+"/PHSP/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
        makeNewKstar(fileName,treeName,True)
        cutFile(fileName,treeName)

#Data
for pol in ["down","up"]:
    for yr in years:
        fileName = mainPath+"/Data/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
        makeNewKstar(fileName,treeName,False)
        cutFile(fileName,treeName)


'''
#Inclusive
for pol in ["down","up"]:
    for yr in years:
        if (yr==2015): continue
        if (yr>2016): continue
        fileName = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
        makeNewKstar(fileName,treeName,True)
        cutFile(fileName,treeName)

for yr in [2011,2012,2016]:
    fileName1 = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+"down/"+str(yr)+"down_pi0Resolved.root"   
    fileName1 = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+"up/"+str(yr)+"up_pi0Resolved.root"
    outname = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/"+str(yr)+"_Inclusive_BDTinput.root"
    mergeFiles(fileName1, fileName1, outname, "DecayTreeTruthMatched")
'''