from collections import namedtuple
import subprocess
import re
import json
import glob
import os
from os.path import normpath, join, isfile
import sys

from Configurables import DaVinci, CondDB, MCDecayTreeTuple

if "CHARMWGPRODROOT" in os.environ:
    # We are running a released version
    sys.path.append(join(os.environ["CHARMWGPRODROOT"], 'productions/d2hll'))
elif "CI_PROJECT_DIR" in os.environ:
    # We are running inside the CI
    sys.path.append(join(os.environ["CI_PROJECT_DIR"], 'productions/d2hll'))
else:
    # We are running with a local version
    sys.path.append(normpath(os.getcwd()))

try:
    from make_decay_tree_tuples import make_decay_tree_tuple
    from decay_parser import parse_decay
    from decay_parser import utils as decay_utils
except ImportError as e:
    raise Exception(str(os.listdir(os.getcwd())) + '\n' + repr(e))

# Keep the linter happy
if False:
    unicode = str


Decay = namedtuple('Decay', ['name', 'decay_descriptor', 'stripping_line', 'dp_event_type', 'dsp_event_type'])

decays = [
    Decay('DpToKmumu_OS', '[D+ -> K+ mu+ mu-]CC', 'D2XMuMuSS_KOSLine', '21113001', '23113001'),
    Decay('DpToKmumu_SS', '[D+ -> K- mu+ mu+]CC', 'D2XMuMuSS_KSSLine', '21113003', '23113004'),
    Decay('DpToKmue_OS', '[D+ -> K+ mu+ e-]CC', 'D2XMuMuSS_K_MuE_OSLine', '21113015', '23113015'),
    Decay('DpToKmue_SS', '[D+ -> K- mu+ e+]CC', 'D2XMuMuSS_K_MuE_SSLine', '21113035', '23113035'),
    Decay('DpToKemu_OS', '[D+ -> K+ e+ mu-]CC', 'D2XMuMuSS_K_EMu_OSLine', '21113005', '23113010'),
    Decay('DpToKee_OS', '[D+ -> K+ e+ e-]CC', 'D2XMuMuSS_K_EE_OSLine', '21123000', '23123000'),
    Decay('DpToKee_SS', '[D+ -> K- e+ e+]CC', 'D2XMuMuSS_K_EE_SSLine', '21123010', '23123010'),
    Decay('DpTopimumu_OS', '[D+ -> pi+ mu+ mu-]CC', 'D2XMuMuSS_PiOSLine', '21113000', '23113000'),
    Decay('DpTopimumu_SS', '[D+ -> pi- mu+ mu+]CC', 'D2XMuMuSS_PiSSLine', '21113002', '23113003'),
    Decay('DpTopimue_OS', '[D+ -> pi+ mu+ e-]CC', 'D2XMuMuSS_Pi_MuE_OSLine', '21113016', '23113016'),
    Decay('DpTopimue_SS', '[D+ -> pi- mu+ e+]CC', 'D2XMuMuSS_Pi_MuE_SSLine', '21113036', '23113036'),
    Decay('DpTopiemu_OS', '[D+ -> pi+ e+ mu-]CC', 'D2XMuMuSS_Pi_EMu_OSLine', '21113006', '23113011'),
    Decay('DpTopiee_OS', '[D+ -> pi+ e+ e-]CC', 'D2XMuMuSS_Pi_EE_OSLine', '21123001', '23123001'),
    Decay('DpTopiee_SS', '[D+ -> pi- e+ e+]CC', 'D2XMuMuSS_Pi_EE_SSLine', '21123011', '23123011'),
    # Unphysical channels
    Decay('DpTopimumu_WS', '[D+ -> pi+ mu+ mu+]CC', 'D2XMuMuSS_PiMuMuCalLine', None, None),
    Decay('DpTopimue_WS', '[D+ -> pi+ mu+ e+]CC', 'D2XMuMuSS_PiEMuCalLine', None, None),
    Decay('DpTopiee_WS', '[D+ -> pi+ e+ e+]CC', 'D2XMuMuSS_PiEECalLine', None, None),
    Decay('DpToKmumu_WS', '[D+ -> K+ mu+ mu+]CC', 'D2XMuMuSS_KMuMuCalLine', None, None),
    Decay('DpToKmue_WS', '[D+ -> K+ mu+ e+]CC', 'D2XMuMuSS_KEMuCalLine', None, None),
    Decay('DpToKee_WS', '[D+ -> K+ e+ e+]CC', 'D2XMuMuSS_KEECalLine', None, None),
    # Calibration channels
    Decay('DpToKKpi_SS', '[D+ -> K- K+ pi+]CC', 'D2XMuMuSS_2KPiLine', None, None),
    Decay('DpToKpipi_SS', '[D+ -> K- pi+ pi+]CC', 'D2XMuMuSS_K2PiLine', None, None),
    Decay('DpTopipipi_SS', '[D+ -> pi- pi+ pi+]CC', 'D2XMuMuSS_PiCalLine', None, None),
]

norms = [
    Decay('DpToKphi_phiTomumu_OS', '[D+ -> K+ mu+ mu-]CC', 'D2XMuMuSS_KOSLine', None, '23113002'),
    Decay('DpTopiphi_phiTomumu_OS', '[D+ -> pi+ mu+ mu-]CC', 'D2XMuMuSS_PiOSLine', '21173001', '23173001'),
    Decay('DpTopiphi_phiToemu_OS', '[D+ -> pi+ e+ mu-]CC', 'D2XMuMuSS_Pi_EMu_OSLine', '21313000', None),
    Decay('DpTopiphi_phiToee_OS', '[D+ -> pi+ e+ e-]CC', 'D2XMuMuSS_Pi_EE_OSLine', '21123020', '23123022'),

    Decay('DspToXphi_phiToemu_OS', '[D+ -> pi+ e+ mu-]CC', 'D2XMuMuSS_Pi_EMu_OSLine', None, '23712012'),
    Decay('DspToXphi_phiTomumu_OS', '[D+ -> pi+ mu+ mu-]CC', 'D2XMuMuSS_PiOSLine', None, '23712020'),
    Decay('DspToXphi_phiToee_OS', '[D+ -> pi+ e+ e-]CC', 'D2XMuMuSS_Pi_EE_OSLine', None, '23722012'),
]


def get_file_info_using_bender():
    file_info_fn = 'file_info_from_bender.json'
    if not isfile(file_info_fn):
        # Load the production configuration so we know what to do
        prod_conf_matches = glob.glob('prodConf_DaVinci_*.py')
        assert len(prod_conf_matches) == 1, prod_conf_matches
        with open(prod_conf_matches[0], 'rt') as fp:
            prod_conf = fp.read()

        # Find the XMLFileCatalog filename
        matches = re.findall("XMLFileCatalog=['\"]([^'\"]+\.xml)['\"]", prod_conf)
        assert len(matches) == 1, (matches, prod_conf)
        xml_catalog_fn = matches[0]

        # Find the LFNs
        lfns = re.findall("['\"](LFN:/[^'\"]+)['\"]", prod_conf)
        assert lfns, (lfns, prod_conf)

        # Use Bender to find the event type, DB tag and data type
        subprocess.check_call(
            "lb-run Bender/latest bender -c '"
            'import json; '
            'from BenderTools.Parser import dataType; '
            'data_type, simulation, input_type = dataType("' + lfns[0] + '"); '
            'json.dump({'
            '"event_type": str(evt["/Event/Gen/Header"].evType()) if simulation else None, '
            '"db_tags": dict(evt["/Event/MC/Header"].condDBTags()) if simulation else None, '
            '"data_type": data_type, '
            '"simulation": simulation, '
            '"input_type": input_type'
            '}, open("file_info_from_bender.json", "wt"))'
            "' --no-color --xml '" + xml_catalog_fn + "' -b "
            "'" + lfns[0] + "'",
            shell=True
        )
    with open('file_info_from_bender.json', 'rt') as fp:
        _file_info = json.load(fp)

    # Convert everything to Python 2 strings to keep Gaudi happy
    file_info = {}
    for k, v in _file_info.items():
        k = str(k)
        if isinstance(v, unicode):
            v = str(v)
        elif isinstance(v, dict):
            v = {str(_k): str(_v) for _k, _v in v.items()}
        file_info[k] = v

    return file_info


file_info = get_file_info_using_bender()

# Finally configure DaVinci
DaVinci().InputType = file_info['input_type']
DaVinci().TupleFile = 'Charm_D2hll_DVntuple.root'
DaVinci().PrintFreq = 100
# DaVinci().EvtMax = 2000
DaVinci().DataType = file_info['data_type']
DaVinci().Simulation = file_info['simulation']
# Only ask for luminosity information when not using simulated data
DaVinci().Lumi = not DaVinci().Simulation

#decay tree location Lb line 61
if DaVinci().Simulation:
    DaVinci().CondDBtag = file_info['db_tags']['SIMCOND']
    DaVinci().DDDBtag = file_info['db_tags']['DDDB']
    if file_info['input_type'] == 'MDST':
        DaVinci().RootInTES = '/Event/AllStreams'
    elif file_info['input_type'] == 'DST':
        pass
    else:
        raise ValueError(file_info)
else:
    db = CondDB(LatestGlobalTagByDataType=file_info['data_type'])
    DaVinci().RootInTES = '/Event/Charm'


# Setup the needed DecayTreeTuples
decays_to_consider = decays[:]

dtts = []
for decay in decays_to_consider:
    dtt_name = decay_descriptor = None

    if DaVinci().Simulation:
        if file_info['event_type'] == decay.dp_event_type:
            dtt_name = decay.name
            decay_descriptor = decay.decay_descriptor
        elif file_info['event_type'] == decay.dsp_event_type:
            dtt_name = decay.name.replace('Dp', 'Dsp')
            decay_descriptor = decay.decay_descriptor.replace('D+', 'D_s+')
        else:
            matches = [norm_decay for norm_decay in norms
                       if file_info['event_type'] == norm_decay.dp_event_type or
                       file_info['event_type'] == norm_decay.dsp_event_type]
            if matches:
                assert len(matches) == 1, matches
                if file_info['event_type'] == matches[0].dp_event_type:
                    dtt_name = decay.name
                    decay_descriptor = decay.decay_descriptor
                elif file_info['event_type'] == matches[0].dsp_event_type:
                    dtt_name = decay.name.replace('Dp', 'Dsp')
                    decay_descriptor = decay.decay_descriptor.replace('D+', 'D_s+')
                else:
                    raise RuntimeError(file_info['event_type'])
            else:
                # Nothing should be added to dtts
                continue
    else:
        dtt_name = decay.name
        decay_descriptor = decay.decay_descriptor

    assert dtt_name and decay_descriptor
    dtts.append(make_decay_tree_tuple(
        name=dtt_name,
        line=decay.stripping_line,
        decay_descriptor=decay_descriptor,
        input_type=file_info['input_type'],
        year=file_info['data_type'],
        mc=DaVinci().Simulation
    ))

    # Add MCDecayTreeTuple so we can get the Generator+Reconstruction+Stripping efficiency
    if DaVinci().Simulation:
        mctuple = MCDecayTreeTuple("MC"+dtt_name)
        assert decay_descriptor.count('->') == 1, decay_descriptor
        _decay = parse_decay(decay_descriptor)
        mctuple.Decay = str(decay_utils.mark_all(_decay)).replace('->', '==>')
        mctuple.addBranches({k: v.replace('->', '==>')
                             for k, v in decay_utils.get_branches(_decay).items()})

        mctuple.ToolList = [
            "MCTupleToolHierarchy",
            "LoKi::Hybrid::MCTupleTool/LoKi_Photos"
        ]
        # Add a 'number of photons' branch
        mctuple.addTupleTool("MCTupleToolKinematic").Verbose = True
        mctuple.addTupleTool("LoKi::Hybrid::TupleTool/LoKi_Photos").Variables = {
            "nPhotos": "MCNINTREE(('gamma' == MCABSID))"
        }
        dtts.append(mctuple)

assert dtts, file_info
DaVinci().UserAlgorithms += dtts