'''
lb-run -c x86_64_v2-el9-clang16-opt DaVinci/v65r0 lbexec DV_advanced1:main MBoptions.yaml
'''
import Functors as F
import FunTuple.functorcollections as FC
from DaVinci import Options, make_config
from DaVinci.algorithms import create_lines_filter
from FunTuple import FunctorCollection
from FunTuple import FunTuple_Particles as Funtuple
from PyConf.reading import get_particles, get_pvs
from RecoConf.event_filters import require_pvs

import Functors as F
from Functors.math import in_range
from RecoConf import standard_particles, algorithms_thor
from RecoConf.reconstruction_objects import make_pvs
    
from GaudiKernel.SystemOfUnits import GeV, MeV, mm, picosecond


def setup_selections():
    make_pions = standard_particles.make_long_pions()
    make_kaons = standard_particles.make_long_kaons()

    ## We need to import the PVs to calculate associated variables
    pvs = make_pvs()

    code_pions = F.require_all(F.PT > 250 * MeV, F.MINIPCHI2(pvs) > 4,)
    code_kaons = F.require_all(F.PT > 250 * MeV, F.MINIPCHI2(pvs) > 4, F.PID_K > 5)

    pions = algorithms_thor.ParticleFilter(make_pions, F.FILTER(code_pions))
    kaons = algorithms_thor.ParticleFilter(make_kaons, F.FILTER(code_kaons))
    
    ds_mass_bounds = in_range(1920 * MeV, F.MASS, 2010 * MeV)

    ds_combination = F.require_all(ds_mass_bounds, F.SDOCA(1, 3) < 0.2 * mm, 
                                F.SDOCA(2, 3) < 0.2 * mm)
    ds_vertex = F.require_all(ds_mass_bounds, F.CHI2DOF < 9)

    ds = algorithms_thor.ParticleCombiner(
        [kaons, kaons, pions], ## The order is important here, must match that of the DecayDescriptor
        name="DsToKpKmPip",
        DecayDescriptor="[D_s- -> K+ K- pi-]cc",
        CombinationCut=ds_combination,
        CompositeCut=ds_vertex)

    b_mass_bounds = in_range(5050 * MeV, F.MASS, 5650 * MeV)

    b_combination = F.require_all(b_mass_bounds)
    b_vertex = F.require_all(b_mass_bounds, F.BPVIPCHI2(pvs) < 16,
                            F.BPVLTIME(pvs) > 0.2 * picosecond)

    bmeson = algorithms_thor.ParticleCombiner(
        [ds, pions],
        name="BToDsmPip",
        DecayDescriptor="[B0 -> D_s- pi+]cc",
        CombinationCut=b_combination,
        CompositeCut=b_vertex
    )
    
    return bmeson
    

def main(options: Options):
    bmeson = setup_selections()
    pvs = get_pvs()

    fields = {
        "Bs": "[B0 -> (D_s- -> K- K+ pi-) pi+]CC", #Note here, we call the parent a B0 even though it's a Bs, because it needs to match the reconstruction in the trigger line
        "Ds": "[B0 -> ^(D_s- -> K- K+ pi-) pi+]CC",
        "Bs_pi": "[B0 -> (D_s- -> K- K+ pi-) ^pi+]CC",
        "Ds_pi": "[B0 -> (D_s- -> K- K+ ^pi-) pi+]CC",
        "Km": "[B0 -> (D_s- -> ^K- K+ pi-) pi+]CC",
        "Kp": "[B0 -> (D_s- -> K- ^K+ pi-) pi+]CC",
    }

    all_vars = FunctorCollection({
        "M": F.MASS,
        "P": F.P,
        "PT": F.PT
    })

    variables = {"ALL": all_vars}

    funtuple = Funtuple(
        name="MyTuple",
        tuple_name="DecayTree",
        fields=fields,
        variables=variables,
        inputs=bmeson,
    )

    return make_config(options, [require_pvs(pvs), funtuple])