# flake8: noqa


from Moore import options, run_reconstruction

from Moore.config import Reconstruction
from PyConf.Algorithms import PrKalmanFilter, PrMatchNN, fromPrMatchTracksV1Tracks
from PyConf.Tools import TrackMasterExtrapolator, PrMCDebugMatchToolNN
from PyConf.application import make_data_with_FetchDataFromFile
from RecoConf.data_from_file import mc_unpackers


import glob

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,
    get_track_checkers,
)
from RecoConf.core_algorithms import make_unique_id_generator
from RecoConf.hlt2_tracking import make_hlt2_tracks, get_global_ut_hits_tool
from RecoConf.hlt1_tracking import (
    make_VeloClusterTrackingSIMD_hits,
    make_PrStorePrUTHits_hits,
    make_PrStoreSciFiHits_hits,
    get_global_materiallocator,
)

"""
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).

author:     Furkan Cetin
date:       10/2023

Moore/run gaudirun.py /work/cetin/LHCb/reco_tuner/moore_options/get_resolution_and_eff_data.py
"""


decay = "test"

options.evt_max = -1

options.ntuple_file = f"data/resolutions_and_effs_{decay}_preselect_ttracks.root"
options.input_type = "ROOT"


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

options.dddb_tag = "dddb-20210617"
options.conddb_tag = "sim-20210617-vc-md100"
options.simulation = True
options.output_level = 3


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_best = make_links_tracks_mcparticles(
        InputTracks=tracks["BestLong"],
        LinksToLHCbIDs=links_to_lhcbids,
    )

    links_to_forward = make_links_tracks_mcparticles(
        InputTracks=tracks["Forward"],
        LinksToLHCbIDs=links_to_lhcbids,
    )

    links_to_velo = make_links_tracks_mcparticles(
        InputTracks=tracks["Velo"],
        LinksToLHCbIDs=links_to_lhcbids,
    )

    links_to_seed = make_links_tracks_mcparticles(
        InputTracks=tracks["Seed"],
        LinksToLHCbIDs=links_to_lhcbids,
    )

    match_debug = PrMatchNN(
        VeloInput=tracks["Velo"]["Pr"],
        SeedInput=tracks["Seed"]["Pr"],
        MatchDebugToolName=PrMCDebugMatchToolNN(
            VeloTracks=tracks["Velo"]["v1"],
            SeedTracks=tracks["Seed"]["v1"],
            VeloTrackLinks=links_to_velo,
            SeedTrackLinks=links_to_seed,
            TrackInfo=make_data_with_FetchDataFromFile(
                "/Event/MC/TrackInfo", "LHCb::MCProperty"
            ),
            MCParticles=mc_unpackers()["MCParticles"],
        ),
        AddUTHitsToolName=get_global_ut_hits_tool(enable=True),
    ).MatchOutput

    match_tracks = {}
    match_tracks["Pr"] = match_debug
    match_tracks["v1"] = fromPrMatchTracksV1Tracks(
        InputTracksLocation=match_debug,
        VeloTracksLocation=tracks["Velo"]["v1"],
        SeedTracksLocation=tracks["Seed"]["v1"],
    ).OutputTracksLocation

    links_to_match = make_links_tracks_mcparticles(
        InputTracks=match_tracks,
        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",
    )
    res_checker_seed = check_track_resolution(
        tracks["Seed"],
        suffix="Seed",
    )

    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",
        match_tracks,
        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"),
    )
    eff_checker_seed = check_tracking_efficiency(
        "Seed",
        tracks["Seed"],
        links_to_seed,
        links_to_lhcbids,
        get_mc_categories("Seed"),
        get_hit_type_mask("Seed"),
    )

    types_and_locations_for_checkers = {
        "Forward": tracks["Forward"],
        "Seed": tracks["Seed"],
        "Match": match_tracks,
    }

    data = []
    data += get_track_checkers(types_and_locations_for_checkers)
    # data = [
    #     res_checker_forward,
    #     res_checker_best_long,
    #     res_checker_best_forward,
    #     res_checker_seed,
    #     eff_checker_forward,
    #     eff_checker_match,
    #     eff_checker_best_long,
    #     eff_checker_seed,
    # ]

    return Reconstruction("run_tracking_debug", data)


run_reconstruction(options, run_tracking_resolution)