import basf2 as b2
import modularAnalysis as ma
import vertex as vtx
from variables import variables as vm
from variables import utils as vu
from variables import collections as vc

main = b2.Path()

ma.inputMdst('/home/belle2/fmeier/BelleIISummerSchool2022/MC15ri_b_ccbar.mdst.root', path=main)

looseSelection = 'dr < 0.5 and abs(dz) < 2 and inCDCAcceptance'

ma.fillParticleList("K+:sel", f"binaryPID(321, 211) > 0.5 and {looseSelection}", path=main)
ma.fillParticleList("pi+:sel", f"pionID > 0.1 and {looseSelection}", path=main)
ma.fillParticleList("pi+:slow", f"p < 0.5 and {looseSelection}", path=main)

ma.reconstructDecay("D0 -> K-:sel pi+:sel", "1.8 < M < 1.9", path=main)
ma.reconstructDecay('D*+ -> D0 pi+:slow', '0.14 < massDifference(0) < 0.15', dmID=1, path=main)

ma.matchMCTruth("D*+", path=main)

ma.variablesToExtraInfo('D*+', {'daughter(0, InvM)': 'Dz_BFInvM'}, path=main)

vtx.treeFit('D*+', conf_level=0, massConstraint=['D0'], ipConstraint=True, path=main)

vm.addAlias('DzInvM', 'extraInfo(Dz_BFInvM)')
vm.addAlias('DeltaM', 'massDifference(0)')

variableList = ['DzInvM', 'DeltaM', 'chiProb']

variableList += vu.create_aliases_for_selected(vc.kinematics + vc.mc_truth,
                                               '^D*+ -> [^D0 -> ^K- ^pi+] ^pi+',
                                               prefix=['Dst', 'Dz', 'K', 'pi', 'pis']
                                               )
variableList += vu.create_aliases_for_selected(vc.pid,
                                               'D*+ -> [D0 -> ^K- ^pi+] ^pi+',
                                               prefix=['K', 'pi', 'pis']
                                               )

ma.variablesToNtuple("D*+", variables=variableList, filename='intermediate_ntuple.root', treename='tree', path=main)

b2.process(main)