from Moore import options, run_reconstruction
from Moore.config import Reconstruction
from PRConfig.TestFileDB import test_file_db
from PyConf.Algorithms import (
    PrForwardTrackingVelo,
    PrForwardTracking,
    PrTrackAssociator,
    PrMatchNN,
)
from PyConf.application import make_data_with_FetchDataFromFile
from PyConf.Tools import PrMCDebugForwardTool, PrMCDebugMatchToolNN
from RecoConf.data_from_file import mc_unpackers
from RecoConf.hlt1_tracking import make_hlt1_tracks, make_PrStoreSciFiHits_hits
from RecoConf.hlt2_tracking import get_global_ut_hits_tool, make_PrHybridSeeding_tracks
from RecoConf.mc_checking import make_links_lhcbids_mcparticles_tracking_system


options.evt_max = -1
n_files_per_cat = 1
polarity = "MU"
options.ntuple_file = f"data/ghost_data_{polarity}.root"
input_files = (
    (
        test_file_db["upgrade_DC19_01_Bs2JPsiPhi_MD"].filenames[:n_files_per_cat]
        if polarity == "MD"
        else test_file_db["upgrade_DC19_01_Bs2JpsiPhiMU"].filenames[:n_files_per_cat]
    )
    + test_file_db[f"upgrade_DC19_01_Bs2PhiPhi{polarity}"].filenames[:n_files_per_cat]
    + test_file_db[f"upgrade_DC19_01_Z2mumu{polarity}"].filenames[:n_files_per_cat]
    + test_file_db[f"upgrade_DC19_01_Bd2Dstmumu{polarity}"].filenames[:n_files_per_cat]
    + test_file_db[f"upgrade_DC19_01_Dst2D0pi{polarity}"].filenames[:n_files_per_cat]
    + test_file_db[f"upgrade_DC19_01_Bd2Kstee{polarity}"].filenames[:n_files_per_cat]
    + test_file_db[f"upgrade_DC19_01_Dp2KSPip_{polarity}"].filenames[:n_files_per_cat]
)
options.input_files = input_files
options.input_type = "ROOT"
options.set_conds_from_testfiledb(f"upgrade_DC19_01_Dst2D0pi{polarity}")


def run_tracking_debug():
    links_to_hits = make_links_lhcbids_mcparticles_tracking_system()
    hlt1_tracks = make_hlt1_tracks()
    seed_tracks = make_PrHybridSeeding_tracks()

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

    links_to_upstream_tracks = PrTrackAssociator(
        SingleContainer=hlt1_tracks["Upstream"]["v1"],
        LinkerLocationID=links_to_hits,
        MCParticleLocation=mc_unpackers()["MCParticles"],
        MCVerticesInput=mc_unpackers()["MCVertices"],
    ).OutputLocation

    links_to_seed_tracks = PrTrackAssociator(
        SingleContainer=seed_tracks["v1"],
        LinkerLocationID=links_to_hits,
        MCParticleLocation=mc_unpackers()["MCParticles"],
        MCVerticesInput=mc_unpackers()["MCVertices"],
    ).OutputLocation

    # be more robust against imperfect data
    loose_forward_params = dict(
        MaxChi2PerDoF=16,
        MaxChi2XProjection=30,
        MaxChi2PerDoFFinal=8,
        MaxChi2Stereo=8,
        MaxChi2StereoAdd=8,
    )

    forward_debug = PrForwardTrackingVelo(
        InputTracks=hlt1_tracks["Velo"]["Pr"],
        SciFiHits=make_PrStoreSciFiHits_hits(),
        AddUTHitsToolName=get_global_ut_hits_tool(enable=True),
        DebugTool=PrMCDebugForwardTool(
            InputTracks=hlt1_tracks["Velo"]["v1"],
            InputTrackLinks=links_to_velo_tracks,
            MCParticles=mc_unpackers()["MCParticles"],
            SciFiHitLinks=links_to_hits,
            SciFiHits=make_PrStoreSciFiHits_hits(),
            TrackInfo=make_data_with_FetchDataFromFile("/Event/MC/TrackInfo"),
        ),
        **loose_forward_params,
    )

    forward_ut_debug = PrForwardTracking(
        SciFiHits=make_PrStoreSciFiHits_hits(),
        InputTracks=hlt1_tracks["Upstream"]["Pr"],
        AddUTHitsToolName=get_global_ut_hits_tool(enable=True),
        DebugTool=PrMCDebugForwardTool(
            InputTracks=hlt1_tracks["Upstream"]["v1"],
            InputTrackLinks=links_to_upstream_tracks,
            MCParticles=mc_unpackers()["MCParticles"],
            SciFiHitLinks=links_to_hits,
            SciFiHits=make_PrStoreSciFiHits_hits(),
            TrackInfo=make_data_with_FetchDataFromFile("/Event/MC/TrackInfo"),
        ),
        **loose_forward_params,
    )

    loose_matching_params = dict(MaxMatchChi2=30.0, MaxDistX=500, MaxDistY=500)

    match_debug = PrMatchNN(
        VeloInput=hlt1_tracks["Velo"]["Pr"],
        SeedInput=seed_tracks["Pr"],
        MatchDebugToolName=PrMCDebugMatchToolNN(
            VeloTracks=hlt1_tracks["Velo"]["v1"],
            SeedTracks=seed_tracks["v1"],
            VeloTrackLinks=links_to_velo_tracks,
            SeedTrackLinks=links_to_seed_tracks,
            TrackInfo=make_data_with_FetchDataFromFile("/Event/MC/TrackInfo"),
            MCParticles=mc_unpackers()["MCParticles"],
        ),
        AddUTHitsToolName=get_global_ut_hits_tool(enable=True),
        **loose_matching_params,
    ).MatchOutput

    data = [forward_debug, forward_ut_debug, match_debug]

    return Reconstruction("run_tracking_debug", data)


run_reconstruction(options, run_tracking_debug)