###############################################################################
# (c) Copyright 2020 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""
Lines to select the decay D*+ -> D0 pi+ with D0 -> KS0 KS0 and KS0 -> pi+ pi-.
The KS0 can be reconstructed either from long, downstream or upstream track, resulting
in the combinations LLLL, LLLD, LLDD, LDDD, DDDD, ULLL, ULDD.

"""

import Functors as F
from Functors.math import in_range
from GaudiKernel.SystemOfUnits import MeV, mm, GeV
from Moore.config import register_line_builder
from Moore.lines import Hlt2Line
from RecoConf.reconstruction_objects import make_pvs
from Hlt2Conf.standard_particles import make_down_pions, make_long_pions, make_up_pions
from Hlt2Conf.algorithms import ParticleContainersMerger
from Hlt2Conf.algorithms_thor import (ParticleCombiner, ParticleFilter)
from Hlt2Conf.lines.charm.particle_properties import _KS_M
from Hlt2Conf.lines.charm.prefilters import charm_prefilters


#####################################################################
### Shortcuts for filters and builders used throughout the module ###
#####################################################################


def make_tagging_pions():
    return ParticleFilter(
        make_long_pions(),
        F.FILTER(F.require_all(F.PT > 200 * MeV, F.P > 1 * GeV)),
    )


def _make_long_pions_from_ks():
    return ParticleFilter(
        make_long_pions(),
        F.FILTER(F.MINIPCHI2CUT(IPChi2Cut=36., Vertices=make_pvs())))


def _make_down_pions_from_ks():
    return ParticleFilter(
        make_down_pions(),
        F.FILTER(F.require_all(F.PT > 175 * MeV, F.P > 3000 * MeV)))


def _make_down_pions_from_ks_LD():
    return ParticleFilter(
        make_down_pions(),
        F.FILTER(F.require_all(F.PT > 100 * MeV, F.P > 2000 * MeV)))


def _make_up_pions_from_ks():
    return ParticleFilter(
        make_up_pions(),
        F.FILTER(F.MINIPCHI2CUT(IPChi2Cut=10., Vertices=make_pvs())))


def _make_ll_ks(descriptor="KS0 -> pi+ pi-"):
    """Returns maker for KS0 -> pi+ pi- constructed with two long pions."""
    return ParticleCombiner(
        [_make_long_pions_from_ks(),
         _make_long_pions_from_ks()],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Ks_LL_{hash}',
        CombinationCut=F.require_all(
            in_range(_KS_M - 50 * MeV, F.MASS, _KS_M + 50 * MeV),
            F.PT > 450 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(_KS_M - 30 * MeV, F.MASS, _KS_M + 30 * MeV),
            F.CHI2DOF < 7,
            in_range(-100 * mm, F.END_VZ, 500 * mm),
            F.PT > 500 * MeV,
        ),
    )


def _make_ll_ks_ULLL(descriptor="KS0 -> pi+ pi-"):
    """Returns maker for KS0 -> pi+ pi- constructed with two long pions only for ULLL."""
    return ParticleCombiner(
        [_make_long_pions_from_ks(),
         _make_long_pions_from_ks()],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Ks_LL_ULLL_{hash}',
        CombinationCut=F.require_all(
            in_range(_KS_M - 50 * MeV, F.MASS, _KS_M + 50 * MeV),
            F.PT > 250 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(_KS_M - 30 * MeV, F.MASS, _KS_M + 30 * MeV),
            F.CHI2DOF < 7,
            in_range(-100 * mm, F.END_VZ, 500 * mm),
            F.PT > 300 * MeV,
        ),
    )


def _make_ld_ks(descriptor="KS0 -> pi+ pi-"):
    """Returns maker for KS0 -> pi+ pi- constructed with a long and downstream pion."""
    return ParticleCombiner(
        [_make_long_pions_from_ks(),
         _make_down_pions_from_ks_LD()],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Ks_LD_{hash}',
        CombinationCut=F.require_all(
            in_range(_KS_M - 85 * MeV, F.MASS, _KS_M + 85 * MeV),
            F.PT > 450 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(_KS_M - 65 * MeV, F.MASS, _KS_M + 65 * MeV),
            F.CHI2DOF < 10,
            in_range(50 * mm, F.END_VZ, 500 * mm),
            F.PT > 500 * MeV,
        ),
    )


def _make_dd_ks(descriptor="KS0 -> pi+ pi-"):
    """Returns maker for KS0 -> pi+ pi- constructed with two downstream pions."""
    return ParticleCombiner(
        [_make_down_pions_from_ks(),
         _make_down_pions_from_ks()],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Ks_DD_{hash}',
        CombinationCut=F.require_all(
            in_range(_KS_M - 80 * MeV, F.MASS, _KS_M + 80 * MeV),
            F.PT > 450 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(_KS_M - 60 * MeV, F.MASS, _KS_M + 60 * MeV),
            F.CHI2DOF < 10,
            in_range(300 * mm, F.END_VZ, 2275 * mm),
            F.PT > 500 * MeV,
        ),
    )


def _make_ul_ks(descriptor="KS0 -> pi+ pi-"):
    """Returns maker for KS0 -> pi+ pi- constructed with one long and one upstream pion."""
    return ParticleCombiner(
        [_make_up_pions_from_ks(),
         _make_long_pions_from_ks()],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Ks_UL_{hash}',
        CombinationCut=F.require_all(
            in_range(_KS_M - 90 * MeV, F.MASS, _KS_M + 90 * MeV),
            F.PT < 5050 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(_KS_M - 70 * MeV, F.MASS, _KS_M + 70 * MeV),
            F.CHI2DOF < 10,
            in_range(-100 * mm, F.END_VZ, 500 * mm),
            F.PT > 400 * MeV,
        ),
    )


def _make_dstars(dzeros, pions, descriptor):
    """Returns maker for D*+- -> D0 pi+-."""
    return ParticleCombiner(
        [dzeros, pions],
        DecayDescriptor=descriptor,
        name='Charm_D0ToKsKs_Dstars_{hash}',
        CombinationCut=F.MASS - F.CHILD(1, F.MASS) < 160 * MeV,
        CompositeCut=F.require_all(
            F.MASS - F.CHILD(1, F.MASS) < 150 * MeV,
            F.CHI2DOF < 25,  # fix due to control channel
        ),
    )


#########################
### Lines definition  ###
#########################

all_lines = {}


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_llll(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_LLLL"):
    pvs = make_pvs()
    kshorts = _make_ll_ks()
    dzeros = ParticleCombiner(
        [kshorts, kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_LLLL",
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.BPVFDCHI2(pvs) > 5,
            F.CHI2DOF < 10,  # fix
            F.BPVDIRA(pvs) > 0.9994,
        ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_llld_line(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_LLLD"):
    pvs = make_pvs()
    ll_kshorts = _make_ll_ks()
    ld_kshorts = _make_ld_ks()
    dzeros = ParticleCombiner(
        [ll_kshorts, ld_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_LLLD",
        AllowDiffInputsForSameIDChildren=True,
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
            ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
            ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])

    return Hlt2Line(name=name, algs=charm_prefilters() + [ld_kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_lldd(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_LLDD"):
    pvs = make_pvs()
    dd_kshorts = _make_dd_ks()
    ll_kshorts = _make_ll_ks()
    dzeros = ParticleCombiner(
        [dd_kshorts, ll_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        AllowDiffInputsForSameIDChildren=True,
        name="Charm_D0ToKsKs_D0ToKsKs_LLDD",
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
        ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [dd_kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_lddd(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_LDDD"):
    pvs = make_pvs()
    dd_kshorts = _make_dd_ks()
    ld_kshorts = _make_ld_ks()
    dzeros = ParticleCombiner(
        [dd_kshorts, ld_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_LDDD",
        AllowDiffInputsForSameIDChildren=True,
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
            ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
            ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [ld_kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_dddd(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_DDDD"):
    pvs = make_pvs()
    dd_kshorts = _make_dd_ks()
    dzeros = ParticleCombiner(
        [dd_kshorts, dd_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_DDDD",
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
        ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [dd_kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_ulll(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_ULLL"):
    pvs = make_pvs()
    ll_kshorts = _make_ll_ks_ULLL()
    ul_kshorts = _make_ul_ks()
    dzeros = ParticleCombiner(
        [ll_kshorts, ul_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_ULLL",
        AllowDiffInputsForSameIDChildren=True,
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
        ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, make_tagging_pions(), "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, make_tagging_pions(), "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [ul_kshorts, dzeros, dstars])


@register_line_builder(all_lines)
def dst_to_d0pi_d0toksks_uldd(
        name="Hlt2Charm_DstpToD0Pip_D0ToKsKs_ULDD"):
    pvs = make_pvs()
    dd_kshorts = _make_dd_ks()
    ul_kshorts = _make_ul_ks()
    dzeros = ParticleCombiner(
        [dd_kshorts, ul_kshorts],
        DecayDescriptor="D0 -> KS0 KS0",
        name="Charm_D0ToKsKs_D0ToKsKs_ULDD",
        AllowDiffInputsForSameIDChildren=True,
        CombinationCut=F.require_all(
            in_range(1730 * MeV, F.MASS, 2000 * MeV),
            F.SUM(F.PT) > 1500 * MeV,
        ),
        CompositeCut=F.require_all(
            in_range(1775 * MeV, F.MASS, 1955 * MeV),
            F.CHI2DOF < 10,
            F.BPVDIRA(pvs) > 0.9994,
        ),
    )

    pitag = make_tagging_pions()

    dstarp = _make_dstars(dzeros, pitag, "D*(2010)+ -> D0 pi+")
    dstarm = _make_dstars(dzeros, pitag, "D*(2010)- -> D0 pi-")
    dstars = ParticleContainersMerger([dstarp, dstarm])


    return Hlt2Line(
        name=name, algs=charm_prefilters() + [ul_kshorts, dzeros, dstars])