"""
This set of options is used for reconstruction development purposes,
and assumes that the input contains MCHits (i.e. is of `Exended`
DST/digi type).
"""
# flake8: noqaq


from Moore import options, run_reconstruction

from Moore.config import Reconstruction
from PyConf.Algorithms import PrKalmanFilter
from PyConf.Tools import TrackMasterExtrapolator

from RecoConf.mc_checking import (
    check_track_resolution,
    check_tracking_efficiency,
    get_mc_categories,
    get_hit_type_mask,
    make_links_lhcbids_mcparticles_tracking_system,
    make_links_tracks_mcparticles,
)
from RecoConf.core_algorithms import make_unique_id_generator
from RecoConf.hlt2_tracking import make_hlt2_tracks
from RecoConf.hlt1_tracking import (
    make_VeloClusterTrackingSIMD_hits,
    make_PrStorePrUTHits_hits,
    make_PrStoreSciFiHits_hits,
    get_global_materiallocator,
)

# sample = "Bd2KstEE_MD"
# sample = "Bd2KstEE_MU"
# sample = "Bs2JpsiPhi_MD"
# sample = "Bs2JpsiPhi_MU"
sample = "Bs2PhiPhi_MD"
# sample = "Bs2PhiPhi_MU"

stack = ""

options.evt_max = 5000
options.first_evt = 0 if stack else 5000
options.ntuple_file = f"data/resolutions_and_effs_{sample}{stack}.root"
options.input_type = "ROOT"
options.set_input_and_conds_from_testfiledb(f"upgrade_sim10_Up08_Digi15_{sample}")


def run_tracking_resolution():
    tracks = make_hlt2_tracks(light_reco=True, fast_reco=False, use_pr_kf=True)
    fitted_forward_tracks = PrKalmanFilter(
        Input=tracks["Forward"]["Pr"],
        MaxChi2=2.8,
        MaxChi2PreOutlierRemoval=20,
        HitsVP=make_VeloClusterTrackingSIMD_hits(),
        HitsUT=make_PrStorePrUTHits_hits(),
        HitsFT=make_PrStoreSciFiHits_hits(),
        ReferenceExtrapolator=TrackMasterExtrapolator(
            MaterialLocator=get_global_materiallocator(),
        ),
        InputUniqueIDGenerator=make_unique_id_generator(),
    ).OutputTracks

    links_to_lhcbids = make_links_lhcbids_mcparticles_tracking_system()
    links_to_forward = make_links_tracks_mcparticles(
        InputTracks=tracks["Forward"],
        LinksToLHCbIDs=links_to_lhcbids,
    )
    links_to_match = make_links_tracks_mcparticles(
        InputTracks=tracks["Match"],
        LinksToLHCbIDs=links_to_lhcbids,
    )
    links_to_best = make_links_tracks_mcparticles(
        InputTracks=tracks["BestLong"],
        LinksToLHCbIDs=links_to_lhcbids,
    )
    res_checker_forward = check_track_resolution(tracks["Forward"], suffix="Forward")
    res_checker_best_long = check_track_resolution(
        tracks["BestLong"],
        suffix="BestLong",
    )
    res_checker_best_forward = check_track_resolution(
        dict(v1=fitted_forward_tracks),
        suffix="BestForward",
    )
    eff_checker_forward = check_tracking_efficiency(
        "Forward",
        tracks["Forward"],
        links_to_forward,
        links_to_lhcbids,
        get_mc_categories("Forward"),
        get_hit_type_mask("Forward"),
    )
    eff_checker_match = check_tracking_efficiency(
        "Match",
        tracks["Match"],
        links_to_match,
        links_to_lhcbids,
        get_mc_categories("Match"),
        get_hit_type_mask("Match"),
    )
    eff_checker_best_long = check_tracking_efficiency(
        "BestLong",
        tracks["BestLong"],
        links_to_best,
        links_to_lhcbids,
        get_mc_categories("BestLong"),
        get_hit_type_mask("BestLong"),
    )
    data = [
        res_checker_forward,
        res_checker_best_long,
        res_checker_best_forward,
        eff_checker_forward,
        eff_checker_match,
        eff_checker_best_long,
    ]

    return Reconstruction("run_tracking_debug", data)


run_reconstruction(options, run_tracking_resolution)
