# flake8: noqa
from Moore import options, run_reconstruction
from RecoConf.hlt2_tracking import (
    make_hlt2_tracks,
    get_default_out_track_types_for_light_reco,
    convert_tracks_to_v3_from_v1,
    get_global_ut_hits_tool,
)
from RecoConf.hlt1_tracking import make_all_pvs
from RecoConf.event_filters import require_gec
from RecoConf.mc_checking import (
    get_track_checkers,
    get_fitted_tracks_checkers,
    check_tracking_efficiency,
    make_links_lhcbids_mcparticles_tracking_system,
    make_links_tracks_mcparticles,
    get_mc_categories,
    get_hit_type_mask,
)
from RecoConf.calorimeter_reconstruction import (
    make_photons_and_electrons,
    make_clusters,
    make_acceptance,
    make_track_cluster_matching,
    make_digits,
    make_track_electron_and_brem_matching,
    make_trackbased_eshower,
)
from Moore.config import Reconstruction
from PyConf.Algorithms import (
    PrFilterTracks2CaloClusters,
    PrMatchNNv3,
    PrFilterTracks2ElectronMatch,
    PrFilterTracks2ElectronShower,
    fromPrMatchTracksV1Tracks,
    fromV3TrackV1Track,
)
import Functors as F
import glob


decay = "BJpsi"

options.evt_max = -1

# options.ntuple_file = f"/work/cetin/LHCb/reco_tuner/data_matching/NewParams/calo_data_{decay}_NewParams.root"  # _dll_NegThree_mlp_NullSix.root"

options.ntuple_file = f"/work/cetin/LHCb/reco_tuner/efficiencies/electrons/calo_effs_{decay}_EcalSelection.root"

# options.ntuple_file = f"data/calo_effs_{decay}.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 == "testJpsi":
    options.input_files = [
        "/auto/data/guenther/Bd_JpsiKst_ee/00143565_00000009_1.xdigi",
        "/auto/data/guenther/Bd_JpsiKst_ee/00143565_00000059_1.xdigi",
        "/auto/data/guenther/Bd_JpsiKst_ee/00143565_00000020_1.xdigi",
    ]
elif decay == "test":
    options.input_files = ["/auto/data/guenther/Bd_Kstee/00151673_00000002_1.xdigi"]

options.input_type = "ROOT"

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


def standalone_hlt2_fastest_reco():
    hlt2_tracks = make_hlt2_tracks(light_reco=True, fast_reco=False, use_pr_kf=True)

    digisEcal = make_digits(calo_raw_bank=False)["digitsEcal"]
    ecalClusters = make_clusters(digisEcal)
    tracks_v3, trackrels = convert_tracks_to_v3_from_v1(
        hlt2_tracks["Seed"]["v1"],
        track_types=["Ttrack"],
    )
    # track acceptances
    tracks_incalo = make_acceptance(tracks_v3)
    tcmatches = make_track_cluster_matching(ecalClusters, tracks_incalo)
    PhElOutput = make_photons_and_electrons(
        ecalClusters, tcmatches["combined"], make_all_pvs()["v3"]
    )
    photons = PhElOutput["photons"]
    electrons = PhElOutput["electrons"]

    eshower = make_trackbased_eshower(tracks_incalo, digisEcal)

    tcmatches_e = make_track_electron_and_brem_matching(
        tracks_incalo, tcmatches, digisEcal, electrons, photons
    )

    # filter with calo clusters
    calo_matched_seeds = PrFilterTracks2CaloClusters(
        Relation=tcmatches["Ttrack"],
        Cut=F.FILTER((F.MIN_ELEMENT_NOTZERO @ F.FORWARDARG0 @ F.WEIGHT) < 20),
    ).Output

    # corrections on track (bit better for elec?)
    # Cut=F.FILTER(F.ALL)   ).Output
    electron_matched_seeds = PrFilterTracks2ElectronMatch(
        Relation=tcmatches_e["Ttrack"]["ElectronMatch"],
        Cut=F.FILTER(F.MIN_ELEMENT_NOTZERO @ F.FORWARDARG0 @ F.WEIGHT < 20),
    ).Output

    # should be best; shape of shower etc; E/p statt chi2; DLL even better? GET(1)
    # Cut=F.FILTER(F.ALL)).Output
    shower_matched_seeds = PrFilterTracks2ElectronShower(
        # Relation=eshower["Ttrack"], Cut=F.FILTER((F.GET(0) @ F.WEIGHT) > 0.7)
        Relation=eshower["Ttrack"],
        Cut=F.FILTER(
            ((F.GET(0) @ F.WEIGHT) > 0.7) | ((F.GET(1) @ F.WEIGHT) > -2)
        ),  # | ((F.GET(0) @ F.WEIGHT) > 0.7)),
    ).Output

    matched_seeds = {}
    matched_seeds["v3"] = shower_matched_seeds
    matched_seeds["v1"] = fromV3TrackV1Track(
        InputTracks=matched_seeds["v3"]
    ).OutputTracks

    matching_params = dict(
        MaxdDist=0.1,
        MinMatchNN=0.25,  # NN response cut value
        MinZMag=5100,
        MaxZMag=5700,
    )

    calo_long = PrMatchNNv3(
        VeloInput=hlt2_tracks["Velo"]["Pr"],
        SeedInput=matched_seeds["v3"],
        AddUTHitsToolName=get_global_ut_hits_tool(),
        **matching_params,
    ).MatchOutput

    match_tracks = {}
    match_tracks["Pr"] = calo_long
    match_tracks["v1"] = fromPrMatchTracksV1Tracks(
        InputTracksLocation=match_tracks["Pr"],
        VeloTracksLocation=hlt2_tracks["Velo"]["v1"],
        SeedTracksLocation=matched_seeds["v1"],
    ).OutputTracksLocation

    out_track_types = get_default_out_track_types_for_light_reco()
    best_tracks = {
        track_type: hlt2_tracks[track_type] for track_type in out_track_types["Best"]
    }

    # data = [
    #     calo_long,
    #     calo_matched_seeds,
    #     electron_matched_seeds,
    #     shower_matched_seeds,
    # ]
    data = []
    types_and_locations_for_checkers = {
        "Forward": hlt2_tracks["Forward"],
        "Match": match_tracks,  # hlt2_tracks["Match"],
        "Seed": hlt2_tracks["Seed"],
    }
    data += get_track_checkers(types_and_locations_for_checkers)
    # data += get_fitted_tracks_checkers(best_tracks)

    return Reconstruction("hlt2_reco", data, [require_gec()])


run_reconstruction(options, standalone_hlt2_fastest_reco)