# Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration

from AthenaConfiguration.ComponentFactory import CompFactory
from TrigBphysHypo.TrigMultiTrkComboHypoMonitoringConfig import TrigMultiTrkComboHypoMonitoring, TrigMultiTrkComboHypoToolMonitoring

from AthenaCommon.Logging import logging
log = logging.getLogger('TrigMultiTrkComboHypoConfig')

trigMultiTrkComboHypoToolDict = {
    'bJpsimumu'     : { 'massRange' : (2500.,  4300.), 'chi2' : 20. },
    'bJpsi'         : { 'massRange' : (2500.,  4300.), 'chi2' : 20. },
    'bJpsimumul2io' : { 'massRange' : (2500.,  4300.), 'chi2' : 20. },
    'bUpsimumu'     : { 'massRange' : (8000., 12000.), 'chi2' : 20. },
    'bUpsi'         : { 'massRange' : (8000., 12000.), 'chi2' : 20. },
    'bDimu'         : { 'massRange' : (1500., 14000.), 'chi2' : 20. },
    'bDimu2700'     : { 'massRange' : ( 100.,  2700.), 'chi2' : 20. },
    'bDimu6000'     : { 'massRange' : ( 100.,  6000.), 'chi2' : 20. },
    'bBmumu'        : { 'massRange' : (4000.,  8500.), 'chi2' : 20. },
    'bPhi'          : { 'massRange' : ( 940.,  1100.), 'chi2' : 10. },
    'bTau'          : { 'massRange' : (   0.,  2700.), 'chi2' : 50. },
}


def StreamerDimuL2ComboHypoCfg(name):
    log.debug('DimuL2ComboHypoCfg.name = %s ', name)

    config = TrigMultiTrkComboHypoConfig()
    hypo = config.ConfigurationComboHypo(
        isStreamer = True,
        trigSequenceName = 'Dimu',
        trigLevel = 'L2',
        trackCollection='HLT_IDTrack_Muon_FTF')
    return hypo

def StreamerDimuL2IOComboHypoCfg(name):
    log.debug('DimuL2IOComboHypoCfg.name = %s ', name)

    config = TrigMultiTrkComboHypoConfig()
    hypo = config.ConfigurationComboHypo(
        isStreamer = True,
        trigSequenceName = 'Dimu',
        trigLevel = 'L2IO')
    return hypo

def DimuEFComboHypoCfg(name):
    log.debug('DimuEFComboHypoCfg.name = %s ', name)

    config = TrigMultiTrkComboHypoConfig()
    hypo = config.ConfigurationComboHypo(
        isStreamer = False,
        trigSequenceName = 'Dimu',
        trigLevel = 'EF',
        outputTrigBphysCollection = 'HLT_DimuEF')
    return hypo

def StreamerDimuEFComboHypoCfg(name):
    log.debug('StreamerDimuEFComboHypoCfg.name = %s ', name)

    config = TrigMultiTrkComboHypoConfig()
    hypo = config.ConfigurationComboHypo(
        isStreamer = True,
        trigSequenceName = 'StreamerDimu',
        trigLevel = 'EF')
    hypo.chi2 = 20.
    hypo.massRanges = [ (100., 6000.) ]
    return hypo

def TrigMultiTrkComboHypoToolFromDict(chainDict):
    config = TrigMultiTrkComboHypoConfig()
    tool = config.ConfigurationComboHypoTool(chainDict)
    return tool

class TrigMultiTrkComboHypoConfig(object):

    def ConfigurationComboHypo(self, isStreamer='False', trigSequenceName='Dimu', trigLevel='L2', trackCollection='', outputTrigBphysCollection='TrigBphysContainer'):

        trigLevelDict = {'L2':0, 'L2IO':1, 'EF':2}

        try:
            value = trigLevelDict[trigLevel]
            log.debug('TrigMultiTrkComboHypo.trigLevel = %s ', value)
        except KeyError:
            raise Exception('TrigMultiTrkComboHypo.trigLevel should be L2, L2IO or EF, but %s provided.', trigLevel)

        baseName = 'Streamer'+trigSequenceName+trigLevel if isStreamer else trigSequenceName+trigLevel

        from TrkExTools.AtlasExtrapolator import AtlasExtrapolator
        VertexFitter = CompFactory.Trk__TrkVKalVrtFitter(
            name = 'TrigBphysFitter_'+baseName,
            FirstMeasuredPoint = False,
            MakeExtendedVertex = False,
            Extrapolator = AtlasExtrapolator())

        VertexPointEstimator = CompFactory.InDet__VertexPointEstimator(
            name = 'VertexPointEstimator_'+baseName,
            MinDeltaR = [-10000., -10000., -10000.],
            MaxDeltaR = [ 10000.,  10000.,  10000.],
            MaxPhi    = [ 10000.,  10000.,  10000.],
            MaxChi2OfVtxEstimation = 2000.)

        tool = CompFactory.TrigMultiTrkComboHypo(
            name = baseName+'ComboHypo',
            isStreamer = isStreamer,
            trigLevel = trigLevel,
            nTracks = 2,
            massRanges = [ (100., 20000.) ],
            TrackCollectionKey = trackCollection,
            TrigBphysCollectionKey = outputTrigBphysCollection,
            VertexFitter = VertexFitter,
            VertexPointEstimator = VertexPointEstimator,
            CheckMultiplicityMap = False,
            MonTool = TrigMultiTrkComboHypoMonitoring('TrigMultiTrkComboHypoMonitoring_'+baseName))

        return tool

    def ConfigurationComboHypoTool(self, chainDict):

        tool = CompFactory.TrigMultiTrkComboHypoTool(chainDict['chainName'])

        try:
            topo = chainDict['topo'][0]
            value = trigMultiTrkComboHypoToolDict[topo]
            tool.massRange = value['massRange']
            tool.chi2 = value['chi2']
            tool.totalCharge = 0
        except LookupError:
            raise Exception('TrigMultiTrkComboHypo misconfigured for \'%s\': topo \'%s\' is not supported.', chainDict['chainName'], topo)

        if 'nocut' in chainDict['topo']:
            tool.AcceptAll = True

        if 'noos' in chainDict['topo']:
            tool.totalCharge = -100 # negative number to indicate no charge cut

        if 'Lxy0' in chainDict['topo']:
            tool.LxyCut = 0.0

        tool.MonTool = TrigMultiTrkComboHypoToolMonitoring('MonTool')
        return tool