###############################################################################
# (c) Copyright 2023 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.                                       #
###############################################################################
# flake8: noqa
from Moore import options, run_reconstruction
from Moore.config import Reconstruction
from RecoConf.data_from_file import mc_unpackers
from RecoConf.hlt1_tracking import make_VeloClusterTrackingSIMD_hits
from RecoConf.hlt2_tracking import (
    make_hlt2_tracks,
    make_PrKalmanFilter_tracks,
    make_PrStorePrUTHits_hits,
    make_PrStoreSciFiHits_hits,
)
from RecoConf.mc_checking import (
    make_links_lhcbids_mcparticles_tracking_system,
    make_links_tracks_mcparticles,
    make_default_IdealStateCreator,
)

from PyConf.Algorithms import PrTrackAssociator, PrDebugTrackingLosses
from PyConf.application import make_data_with_FetchDataFromFile
import glob


"""

run with
Moore/run gaudirun.py /work/cetin/LHCb/reco_tuner/moore_options/get_tracking_losses.py
tested by mc_matching_example.py

"""

decay = "D"

if decay == "B":
    options.input_files = glob.glob("/auto/data/guenther/Bd_Kstee/*.xdigi")
elif decay == "D":
    options.input_files = glob.glob("/auto/data/guenther/Dst_D0ee/*.xdigi")
elif decay == "test":
    options.input_files = ["/auto/data/guenther/Bd_Kstee/00151673_00000002_1.xdigi"]


options.conddb_tag = "sim-20210617-vc-md100"
options.dddb_tag = "dddb-20210617"
options.simulation = True
options.input_type = "ROOT"

options.ntuple_file = f"data/tracking_losses_ntuple_{decay}_match_electron_weights.root"


options.evt_max = -1


def run_tracking_losses():
    links_to_hits = make_links_lhcbids_mcparticles_tracking_system()
    hlt2_tracks = make_hlt2_tracks(light_reco=True, fast_reco=False, use_pr_kf=True)
    vp_hits = make_VeloClusterTrackingSIMD_hits()
    ut_hits = make_PrStorePrUTHits_hits()
    ft_hits = make_PrStoreSciFiHits_hits()
    fitted_match_tracks = make_PrKalmanFilter_tracks(  # fitted_forward_tracks
        input_tracks=hlt2_tracks["Match"]["Pr"],  # Forward
        hits_vp=vp_hits,
        hits_ut=ut_hits,
        hits_ft=ft_hits,
    )

    # add MCLinking to the (fitted) V1 tracks
    links_to_velo_tracks = PrTrackAssociator(
        SingleContainer=hlt2_tracks["Velo"]["v1"],
        LinkerLocationID=links_to_hits,
        MCParticleLocation=mc_unpackers()["MCParticles"],
        MCVerticesInput=mc_unpackers()["MCVertices"],
    ).OutputLocation

    links_to_long_tracks = PrTrackAssociator(
        SingleContainer=hlt2_tracks["Match"]["v1"],  # Forward
        LinkerLocationID=links_to_hits,
        MCParticleLocation=mc_unpackers()["MCParticles"],
        MCVerticesInput=mc_unpackers()["MCVertices"],
    ).OutputLocation

    with PrTrackAssociator.bind(FractionOK=0.5):
        loose_links_to_long_tracks = PrTrackAssociator(
            SingleContainer=hlt2_tracks["Match"]["v1"],  # Forward
            LinkerLocationID=links_to_hits,
            MCParticleLocation=mc_unpackers()["MCParticles"],
            MCVerticesInput=mc_unpackers()["MCVertices"],
        ).OutputLocation

    links_to_fitted_tracks = PrTrackAssociator(
        SingleContainer=fitted_match_tracks,  # fitted_forward_tracks
        LinkerLocationID=links_to_hits,
        MCParticleLocation=mc_unpackers()["MCParticles"],
        MCVerticesInput=mc_unpackers()["MCVertices"],
    ).OutputLocation

    tracking_losses = PrDebugTrackingLosses(
        name="PrDebugTrackingLosses",
        TrackType="Long",
        StudyTracks=hlt2_tracks["Match"]["v1"],  # Forward
        VeloTracks=hlt2_tracks["Velo"]["v1"],
        MCParticles=mc_unpackers()["MCParticles"],
        MCVPHits=mc_unpackers()["MCVPHits"],
        MCUTHits=mc_unpackers()["MCUTHits"],
        MCFTHits=mc_unpackers()["MCFTHits"],
        VeloTrackLinks=links_to_velo_tracks,
        TrackLinks=links_to_long_tracks,
        LooseTrackLinks=loose_links_to_long_tracks,
        FittedTrackLinks=links_to_fitted_tracks,
        # LHCbIDLinks=links_to_hits,
        IdealStateCreator=make_default_IdealStateCreator(),
        TrackInfo=make_data_with_FetchDataFromFile("/Event/MC/TrackInfo"),
    )

    data = [tracking_losses]
    return Reconstruction("run_tracking_losses", data)


run_reconstruction(options, run_tracking_losses)