from basf2 import B2INFO, B2FATAL, B2WARNING
import basf2
import basf2_mva
import inspect
import modularAnalysis as ma
import variables
from variables import utils
import os
import glob
import b2bii
import collections

def getBelleOrBelle2():
    """
    Gets the global ModeCode.
    """
    if b2bii.isB2BII():
        return 'Belle'
    else:
        return 'Belle2'

def setInteractionWithDatabase(downloadFromDatabaseIfNotFound=False, uploadToDatabaseAfterTraining=False):
    """
    Sets the interaction with the database: download trained weight files or upload weight files after training.
    """

    global downloadFlag
    global uploadFlag

    downloadFlag = downloadFromDatabaseIfNotFound
    uploadFlag = uploadToDatabaseAfterTraining

# Default list of aliases that should be used to save the flavor tagging information using VariablesToNtuple
flavor_tagging = ['FBDT_qrCombined', 'FANN_qrCombined', 'qrMC', 'mcFlavorOfOtherB', 'qrGNN',
                  'qpElectron', 'hasTrueTargetElectron', 'isRightCategoryElectron',
                  'qpIntermediateElectron', 'hasTrueTargetIntermediateElectron', 'isRightCategoryIntermediateElectron',
                  'qpMuon', 'hasTrueTargetMuon', 'isRightCategoryMuon',
                  'qpIntermediateMuon', 'hasTrueTargetIntermediateMuon', 'isRightCategoryIntermediateMuon',
                  'qpKinLepton', 'hasTrueTargetKinLepton', 'isRightCategoryKinLepton',
                  'qpIntermediateKinLepton', 'hasTrueTargetIntermediateKinLepton', 'isRightCategoryIntermediateKinLepton',
                  'qpKaon', 'hasTrueTargetKaon', 'isRightCategoryKaon',
                  'qpSlowPion', 'hasTrueTargetSlowPion', 'isRightCategorySlowPion',
                  'qpFastHadron', 'hasTrueTargetFastHadron', 'isRightCategoryFastHadron',
                  'qpLambda', 'hasTrueTargetLambda', 'isRightCategoryLambda',
                  'qpFSC', 'hasTrueTargetFSC', 'isRightCategoryFSC',
                  'qpMaximumPstar', 'hasTrueTargetMaximumPstar', 'isRightCategoryMaximumPstar',
                  'qpKaonPion', 'hasTrueTargetKaonPion', 'isRightCategoryKaonPion']

def add_default_FlavorTagger_aliases():
    """
    This function adds the default aliases for flavor tagging variables
    and defines the collection of flavor tagging variables.
    """

    variables.variables.addAlias('FBDT_qrCombined', 'qrOutput(FBDT)')
    variables.variables.addAlias('FANN_qrCombined', 'qrOutput(FANN)')
    variables.variables.addAlias('qrMC', 'isRelatedRestOfEventB0Flavor')

    variables.variables.addAlias('qrGNN', 'extraInfo(qrGNN)')

    for iCategory in AvailableCategories:
        aliasForQp = 'qp' + iCategory
        aliasForTrueTarget = 'hasTrueTarget' + iCategory
        aliasForIsRightCategory = 'isRightCategory' + iCategory
        variables.variables.addAlias(aliasForQp, 'qpCategory(' + iCategory + ')')
        variables.variables.addAlias(aliasForTrueTarget, 'hasTrueTargets(' + iCategory + ')')
        variables.variables.addAlias(aliasForIsRightCategory, 'isTrueFTCategory(' + iCategory + ')')

    utils.add_collection(flavor_tagging, 'flavor_tagging')

def set_FlavorTagger_pid_aliases():
    """
    This function adds the pid aliases needed by the flavor tagger.
    """
    variables.variables.addAlias('eid_TOP', 'pidPairProbabilityExpert(11, 211, TOP)')
    variables.variables.addAlias('eid_ARICH', 'pidPairProbabilityExpert(11, 211, ARICH)')
    variables.variables.addAlias('eid_ECL', 'pidPairProbabilityExpert(11, 211, ECL)')

    variables.variables.addAlias('muid_TOP', 'pidPairProbabilityExpert(13, 211, TOP)')
    variables.variables.addAlias('muid_ARICH', 'pidPairProbabilityExpert(13, 211, ARICH)')
    variables.variables.addAlias('muid_KLM', 'pidPairProbabilityExpert(13, 211, KLM)')

    variables.variables.addAlias('piid_TOP', 'pidPairProbabilityExpert(211, 321, TOP)')
    variables.variables.addAlias('piid_ARICH', 'pidPairProbabilityExpert(211, 321, ARICH)')

    variables.variables.addAlias('Kid_TOP', 'pidPairProbabilityExpert(321, 211, TOP)')
    variables.variables.addAlias('Kid_ARICH', 'pidPairProbabilityExpert(321, 211, ARICH)')

    if getBelleOrBelle2() == "Belle":
        variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC, SVD), 0.5)')
        variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC, SVD), 0.5)')
        variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC, SVD), 0.5)')
        variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC, SVD), 0.5)')
        variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC, SVD), 0.5)')
    else:
        variables.variables.addAlias('eid_dEdx', 'pidPairProbabilityExpert(11, 211, CDC)')
        variables.variables.addAlias('muid_dEdx', 'pidPairProbabilityExpert(13, 211, CDC)')
        variables.variables.addAlias('piid_dEdx', 'pidPairProbabilityExpert(211, 321, CDC)')
        variables.variables.addAlias('pi_vs_edEdxid', 'pidPairProbabilityExpert(211, 11, CDC)')
        variables.variables.addAlias('Kid_dEdx', 'pidPairProbabilityExpert(321, 211, CDC)')

def set_FlavorTagger_pid_aliases_legacy():
    """
    This function adds the pid aliases needed by the flavor tagger trained for MC13.
    """
    variables.variables.addAlias('eid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, TOP), 0.5)')
    variables.variables.addAlias('eid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, ARICH), 0.5)')
    variables.variables.addAlias('eid_ECL', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, ECL), 0.5)')

    variables.variables.addAlias('muid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, TOP), 0.5)')
    variables.variables.addAlias('muid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, ARICH), 0.5)')
    variables.variables.addAlias('muid_KLM', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, KLM), 0.5)')

    variables.variables.addAlias('piid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, TOP), 0.5)')
    variables.variables.addAlias('piid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, ARICH), 0.5)')

    variables.variables.addAlias('Kid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, TOP), 0.5)')
    variables.variables.addAlias('Kid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, ARICH), 0.5)')

    if getBelleOrBelle2() == "Belle":
        variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC, SVD), 0.5)')
        variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC, SVD), 0.5)')
        variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC, SVD), 0.5)')
        variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC, SVD), 0.5)')
        variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC, SVD), 0.5)')
    else:
        # Removed SVD PID for Belle II MC and data as it is absent in release 4.
        variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC), 0.5)')
        variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC), 0.5)')
        variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC), 0.5)')
        variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC), 0.5)')
        variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC), 0.5)')

def set_GNNFlavorTagger_aliases(categories):
    """
    This function adds aliases for the GNN-based flavor tagger.
    """

    # will be used for target variable 0:B0bar, 1:B0
    variables.variables.addAlias('qrCombined_bit', '(qrCombined+1)/2')
    alias_list = ['qrCombined_bit']

    var_dict = {
        # position
        'dx': 'dx',
        'dy': 'dy',
        'dz': 'dz',
        # mask
        'E': 'E',
        # charge,
        'charge': 'charge',
        # feature
        'px_c': 'px*charge',
        'py_c': 'py*charge',
        'pz_c': 'pz*charge',
        'electronID_c': 'electronID*charge',
        'muonID_c': 'muonID*charge',
        'pionID_c': 'pionID*charge',
        'kaonID_c': 'kaonID*charge',
        'protonID_c': 'protonID*charge',
        'deuteronID_c': 'deuteronID*charge',
        'electronID_noSVD_noTOP_c': 'electronID_noSVD_noTOP*charge',
    }

    # 16 charged particles are used at most
    for rank in range(1, 17):

        for cat in categories:
            listName = AvailableCategories[cat].particleList
            varName = f'QpTrack({listName}, isRightCategory({cat}), isRightCategory({cat}))'

            varWithRank = f'ifNANgiveX(getVariableByRank(pi+:inRoe, FT_p, {varName}, {rank}), 0)'
            aliasWithRank = f'{cat}_rank{rank}'

            variables.variables.addAlias(aliasWithRank, varWithRank)
            alias_list.append(aliasWithRank)

        for alias, var in var_dict.items():
            varWithRank = f'ifNANgiveX(getVariableByRank(pi+:inRoe, FT_p, {var}, {rank}), 0)'
            aliasWithRank = f'{alias}_rank{rank}'

            variables.variables.addAlias(aliasWithRank, varWithRank)
            alias_list.append(aliasWithRank)

    return alias_list

def setInputVariablesWithMask(maskName='all'):
    """
    Set aliases for input variables with ROE mask.
    """
    variables.variables.addAlias('pMissTag_withMask', 'pMissTag('+maskName+')')
    variables.variables.addAlias('cosTPTO_withMask', 'cosTPTO('+maskName+')')
    variables.variables.addAlias('ptTracksRoe_withMask', 'ptTracksRoe('+maskName+')')
    variables.variables.addAlias('pt2TracksRoe_withMask', 'pt2TracksRoe('+maskName+')')
    variables.variables.addAlias('ptTracksRoe_withMask', 'ptTracksRoe('+maskName+')')

def getFastBDTCategories():
    '''
    Helper function for getting the FastBDT categories.
    It's necessary for removing top-level ROOT imports.
    '''
    fastBDTCategories = basf2_mva.FastBDTOptions()
    fastBDTCategories.m_nTrees = 500
    fastBDTCategories.m_nCuts = 8
    fastBDTCategories.m_nLevels = 3
    fastBDTCategories.m_shrinkage = 0.10
    fastBDTCategories.m_randRatio = 0.5
    return fastBDTCategories

def getFastBDTCombiner():
    '''
    Helper function for getting the FastBDT combiner.
    It's necessary for removing top-level ROOT imports.
    '''
    fastBDTCombiner = basf2_mva.FastBDTOptions()
    fastBDTCombiner.m_nTrees = 500
    fastBDTCombiner.m_nCuts = 8
    fastBDTCombiner.m_nLevels = 3
    fastBDTCombiner.m_shrinkage = 0.10
    fastBDTCombiner.m_randRatio = 0.5
    return fastBDTCombiner

def getMlpFANNCombiner():
    '''
    Helper function for getting the MLP FANN combiner.
    It's necessary for removing top-level ROOT imports.
    '''
    mlpFANNCombiner = basf2_mva.FANNOptions()
    mlpFANNCombiner.m_max_epochs = 10000
    mlpFANNCombiner.m_hidden_layers_architecture = "3*N"
    mlpFANNCombiner.m_hidden_activiation_function = "FANN_SIGMOID_SYMMETRIC"
    mlpFANNCombiner.m_output_activiation_function = "FANN_SIGMOID_SYMMETRIC"
    mlpFANNCombiner.m_error_function = "FANN_ERRORFUNC_LINEAR"
    mlpFANNCombiner.m_training_method = "FANN_TRAIN_RPROP"
    mlpFANNCombiner.m_validation_fraction = 0.5
    mlpFANNCombiner.m_random_seeds = 10
    mlpFANNCombiner.m_test_rate = 500
    mlpFANNCombiner.m_number_of_threads = 8
    mlpFANNCombiner.m_scale_features = True
    mlpFANNCombiner.m_scale_target = False
    # mlpFANNCombiner.m_scale_target = True
    return mlpFANNCombiner

# SignalFraction: FBDT feature
# For smooth output set to -1, this will break the calibration.
# For correct calibration set to -2, leads to peaky combiner output.
signalFraction = -2

# Maximal number of events to train each method
maxEventsNumber = 0  # 0 takes all the sampled events. The number in the past was 500000

FTCategoryParameters = collections.namedtuple('FTCategoryParameters',
                                              ['particleList', 'trackName', 'eventName', 'variableName', 'code'])
# Definition of all available categories, 'standard category name':
# ['ParticleList', 'trackLevel category name', 'eventLevel category name',
# 'combinerLevel variable name', 'category code']
AvailableCategories = {
    'Electron':
    FTCategoryParameters('e+:inRoe', 'Electron', 'Electron',
                         'QpOf(e+:inRoe, isRightCategory(Electron), isRightCategory(Electron))',
                         0),
    'IntermediateElectron':
    FTCategoryParameters('e+:inRoe', 'IntermediateElectron', 'IntermediateElectron',
                         'QpOf(e+:inRoe, isRightCategory(IntermediateElectron), isRightCategory(IntermediateElectron))',
                         1),
    'Muon':
    FTCategoryParameters('mu+:inRoe', 'Muon', 'Muon',
                         'QpOf(mu+:inRoe, isRightCategory(Muon), isRightCategory(Muon))',
                         2),
    'IntermediateMuon':
    FTCategoryParameters('mu+:inRoe', 'IntermediateMuon', 'IntermediateMuon',
                         'QpOf(mu+:inRoe, isRightCategory(IntermediateMuon), isRightCategory(IntermediateMuon))',
                         3),
    'KinLepton':
    FTCategoryParameters('mu+:inRoe', 'KinLepton', 'KinLepton',
                         'QpOf(mu+:inRoe, isRightCategory(KinLepton), isRightCategory(KinLepton))',
                         4),
    'IntermediateKinLepton':
    FTCategoryParameters('mu+:inRoe', 'IntermediateKinLepton', 'IntermediateKinLepton',
                         'QpOf(mu+:inRoe, isRightCategory(IntermediateKinLepton), isRightCategory(IntermediateKinLepton))',
                         5),
    'Kaon':
    FTCategoryParameters('K+:inRoe', 'Kaon', 'Kaon',
                         'weightedQpOf(K+:inRoe, isRightCategory(Kaon), isRightCategory(Kaon))',
                         6),
    'SlowPion':
    FTCategoryParameters('pi+:inRoe', 'SlowPion', 'SlowPion',
                         'QpOf(pi+:inRoe, isRightCategory(SlowPion), isRightCategory(SlowPion))',
                         7),
    'FastHadron':
    FTCategoryParameters('pi+:inRoe', 'FastHadron', 'FastHadron',
                         'QpOf(pi+:inRoe, isRightCategory(FastHadron), isRightCategory(FastHadron))',
                         8),
    'Lambda':
    FTCategoryParameters('Lambda0:inRoe', 'Lambda', 'Lambda',
                         'weightedQpOf(Lambda0:inRoe, isRightCategory(Lambda), isRightCategory(Lambda))',
                         9),
    'FSC':
    FTCategoryParameters('pi+:inRoe', 'SlowPion', 'FSC',
                         'QpOf(pi+:inRoe, isRightCategory(FSC), isRightCategory(SlowPion))',
                         10),
    'MaximumPstar':
    FTCategoryParameters('pi+:inRoe', 'MaximumPstar', 'MaximumPstar',
                         'QpOf(pi+:inRoe, isRightCategory(MaximumPstar), isRightCategory(MaximumPstar))',
                         11),
    'KaonPion':
    FTCategoryParameters('K+:inRoe', 'Kaon', 'KaonPion',
                         'QpOf(K+:inRoe, isRightCategory(KaonPion), isRightCategory(Kaon))',
                         12),
}

def getTrainingVariables(category=None):
    """
    Helper function to get training variables.

    NOTE: This function is not called the Expert mode. It is not necessary to be consistent with variables list of weight files.
    """

    KId = {'Belle': 'ifNANgiveX(atcPIDBelle(3,2), 0.5)', 'Belle2': 'kaonID'}
    muId = {'Belle': 'muIDBelle', 'Belle2': 'muonID'}
    eId = {'Belle': 'eIDBelle', 'Belle2': 'electronID'}

    variables = []
    if category == 'Electron' or category == 'IntermediateElectron':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'p',
                     'pt',
                     'cosTheta',
                     eId[getBelleOrBelle2()],
                     'eid_TOP',
                     'eid_ARICH',
                     'eid_ECL',
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(pMissCMS)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'BtagToWBosonVariables(EW90)',
                     'cosTPTO',
                     'chiProb',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('eid_dEdx')
            variables.append('ImpactXY')
            variables.append('distance')

    elif category == 'Muon' or category == 'IntermediateMuon':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'p',
                     'pt',
                     'cosTheta',
                     muId[getBelleOrBelle2()],
                     'muid_TOP',
                     'muid_ARICH',
                     'muid_KLM',
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(pMissCMS)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'BtagToWBosonVariables(EW90)',
                     'cosTPTO',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('muid_dEdx')
            variables.append('ImpactXY')
            variables.append('distance')
            variables.append('chiProb')

    elif category == 'KinLepton' or category == 'IntermediateKinLepton':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'p',
                     'pt',
                     'cosTheta',
                     muId[getBelleOrBelle2()],
                     'muid_TOP',
                     'muid_ARICH',
                     'muid_KLM',
                     eId[getBelleOrBelle2()],
                     'eid_TOP',
                     'eid_ARICH',
                     'eid_ECL',
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(pMissCMS)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'BtagToWBosonVariables(EW90)',
                     'cosTPTO',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('eid_dEdx')
            variables.append('muid_dEdx')
            variables.append('ImpactXY')
            variables.append('distance')
            variables.append('chiProb')

    elif category == 'Kaon':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'p',
                     'pt',
                     'cosTheta',
                     KId[getBelleOrBelle2()],
                     'Kid_dEdx',
                     'Kid_TOP',
                     'Kid_ARICH',
                     'NumberOfKShortsInRoe',
                     'ptTracksRoe',
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(pMissCMS)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'BtagToWBosonVariables(EW90)',
                     'cosTPTO',
                     'chiProb',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('ImpactXY')
            variables.append('distance')

    elif category == 'SlowPion':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'cosTheta',
                     'p',
                     'pt',
                     'pionID',
                     'piid_TOP',
                     'piid_ARICH',
                     'pi_vs_edEdxid',
                     KId[getBelleOrBelle2()],
                     'Kid_dEdx',
                     'Kid_TOP',
                     'Kid_ARICH',
                     'NumberOfKShortsInRoe',
                     'ptTracksRoe',
                     eId[getBelleOrBelle2()],
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(EW90)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'BtagToWBosonVariables(pMissCMS)',
                     'cosTPTO',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('piid_dEdx')
            variables.append('ImpactXY')
            variables.append('distance')
            variables.append('chiProb')

    elif category == 'FastHadron':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'cosTheta',
                     'p',
                     'pt',
                     'pionID',
                     'piid_dEdx',
                     'piid_TOP',
                     'piid_ARICH',
                     'pi_vs_edEdxid',
                     KId[getBelleOrBelle2()],
                     'Kid_dEdx',
                     'Kid_TOP',
                     'Kid_ARICH',
                     'NumberOfKShortsInRoe',
                     'ptTracksRoe',
                     eId[getBelleOrBelle2()],
                     'BtagToWBosonVariables(recoilMassSqrd)',
                     'BtagToWBosonVariables(EW90)',
                     'BtagToWBosonVariables(cosThetaMissCMS)',
                     'cosTPTO',
                     ]
        if getBelleOrBelle2() == "Belle":
            variables.append('BtagToWBosonVariables(pMissCMS)')
            variables.append('ImpactXY')
            variables.append('distance')
            variables.append('chiProb')

    elif category == 'Lambda':
        variables = ['lambdaFlavor',
                     'NumberOfKShortsInRoe',
                     'M',
                     'cosAngleBetweenMomentumAndVertexVector',
                     'lambdaZError',
                     'daughter(0,p)',
                     'daughter(0,useCMSFrame(p))',
                     'daughter(1,p)',
                     'daughter(1,useCMSFrame(p))',
                     'useCMSFrame(p)',
                     'p',
                     'chiProb',
                     ]
        if getBelleOrBelle2() == "Belle2":
            variables.append('daughter(1,protonID)')  # protonID always 0 in B2BII check in future
            variables.append('daughter(0,pionID)')  # not very powerful in B2BII
        else:
            variables.append('distance')

    elif category == 'MaximumPstar':
        variables = ['useCMSFrame(p)',
                     'useCMSFrame(pt)',
                     'p',
                     'pt',
                     'cosTPTO',
                     ]
        if getBelleOrBelle2() == "Belle2":
            variables.append('ImpactXY')
            variables.append('distance')

    elif category == 'FSC':
        variables = ['useCMSFrame(p)',
                     'cosTPTO',
                     KId[getBelleOrBelle2()],
                     'FSCVariables(pFastCMS)',
                     'FSCVariables(cosSlowFast)',
                     'FSCVariables(cosTPTOFast)',
                     'FSCVariables(SlowFastHaveOpositeCharges)',
                     ]
    elif category == 'KaonPion':
        variables = ['extraInfo(isRightCategory(Kaon))',
                     'HighestProbInCat(pi+:inRoe, isRightCategory(SlowPion))',
                     'KaonPionVariables(cosKaonPion)',
                     'KaonPionVariables(HaveOpositeCharges)',
                     KId[getBelleOrBelle2()]
                     ]

    return variables

def FillParticleLists(maskName='all', categories=None, path=None):
    """
    Fills the particle Lists for all categories.
    """

    from vertex import kFit
    readyParticleLists = []

    if categories is None:
        categories = []

    trackCut = 'isInRestOfEvent > 0.5 and passesROEMask(' + maskName + ') > 0.5 and p >= 0'

    for category in categories:
        particleList = AvailableCategories[category].particleList

        if particleList in readyParticleLists:
            continue

        # Select particles in ROE for different categories according to mass hypothesis.
        if particleList == 'Lambda0:inRoe':
            if 'pi+:inRoe' not in readyParticleLists:
                ma.fillParticleList('pi+:inRoe', trackCut, path=path)
                readyParticleLists.append('pi+:inRoe')

            ma.fillParticleList('p+:inRoe', trackCut, path=path)
            ma.reconstructDecay(particleList + ' -> pi-:inRoe p+:inRoe', '1.00<=M<=1.23', False, path=path)
            kFit(particleList, 0.01, path=path)
            ma.matchMCTruth(particleList, path=path)
            readyParticleLists.append(particleList)

        else:
            # Filling particle list for actual category
            ma.fillParticleList(particleList, trackCut, path=path)
            readyParticleLists.append(particleList)

    # Additional particleLists for K_S0
    if getBelleOrBelle2() == 'Belle':
        ma.cutAndCopyList('K_S0:inRoe', 'K_S0:mdst', 'extraInfo(ksnbStandard) == 1 and isInRestOfEvent == 1', path=path)
    else:
        if 'pi+:inRoe' not in readyParticleLists:
            ma.fillParticleList('pi+:inRoe', trackCut, path=path)
        ma.reconstructDecay('K_S0:inRoe -> pi+:inRoe pi-:inRoe', '0.40<=M<=0.60', False, path=path)
        kFit('K_S0:inRoe', 0.01, path=path)

    # Apply BDT-based LID
    if getBelleOrBelle2() == 'Belle2':
        default_list_for_lid_BDT = ['e+:inRoe', 'mu+:inRoe']
        list_for_lid_BDT = []

        for particleList in default_list_for_lid_BDT:
            if particleList in readyParticleLists:
                list_for_lid_BDT.append(particleList)

        if list_for_lid_BDT:  # empty check
            ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
                                  trainingMode=0,  # binary
                                  binaryHypoPDGCodes=(11, 211))  # e vs pi
            ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
                                  trainingMode=0,  # binary
                                  binaryHypoPDGCodes=(13, 211))  # mu vs pi
            ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
                                  trainingMode=1)  # Multiclass

def eventLevel(mode='Expert', weightFiles='B2JpsiKs_mu', categories=None, path=None):
    """
    Samples data for training or tests all categories all categories at event level.
    """

    from basf2 import create_path
    from basf2 import register_module

    B2INFO('EVENT LEVEL')

    ReadyMethods = 0

    # Each category has its own Path in order to be skipped if the corresponding particle list is empty
    identifiersExtraInfosDict = dict()
    identifiersExtraInfosKaonPion = []

    if categories is None:
        categories = []

    for category in categories:
        particleList = AvailableCategories[category].particleList

        methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
        identifierEventLevel = methodPrefixEventLevel
        targetVariable = 'isRightCategory(' + category + ')'
        extraInfoName = targetVariable

        if mode == 'Expert':

            if downloadFlag or useOnlyLocalFlag:
                identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'

            if downloadFlag:
                if not os.path.isfile(identifierEventLevel):
                    basf2_mva.download(methodPrefixEventLevel, identifierEventLevel)
                    if not os.path.isfile(identifierEventLevel):
                        B2FATAL('Flavor Tagger: Weight file ' + identifierEventLevel +
                                ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')

            if useOnlyLocalFlag:
                if not os.path.isfile(identifierEventLevel):
                    B2FATAL('Flavor Tagger: ' + particleList + ' Eventlevel was not trained. Weight file ' +
                            identifierEventLevel + ' was not found. Stopped')

            B2INFO('flavorTagger: MVAExpert ' + methodPrefixEventLevel + ' ready.')

        elif mode == 'Sampler':

            identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'
            if os.path.isfile(identifierEventLevel):
                B2INFO('flavorTagger: MVAExpert ' + methodPrefixEventLevel + ' ready.')

                if 'KaonPion' in categories:
                    methodPrefixEventLevelKaonPion = "FlavorTagger_" + getBelleOrBelle2() + \
                        "_" + weightFiles + 'EventLevelKaonPionFBDT'
                    identifierEventLevelKaonPion = filesDirectory + '/' + methodPrefixEventLevelKaonPion + '_1.root'
                    if not os.path.isfile(identifierEventLevelKaonPion):
                        # Slow Pion and Kaon categories are used if Kaon-Pion is lacking for
                        # sampling. The others are not needed and skipped
                        if category != "SlowPion" and category != "Kaon":
                            continue

        if mode == 'Expert' or (mode == 'Sampler' and os.path.isfile(identifierEventLevel)):

            B2INFO('flavorTagger: Applying MVAExpert ' + methodPrefixEventLevel + '.')

            if category == 'KaonPion':
                identifiersExtraInfosKaonPion.append((extraInfoName, identifierEventLevel))
            elif particleList not in identifiersExtraInfosDict:
                identifiersExtraInfosDict[particleList] = [(extraInfoName, identifierEventLevel)]
            else:
                identifiersExtraInfosDict[particleList].append((extraInfoName, identifierEventLevel))

            ReadyMethods += 1

    # Each category has its own Path in order to be skipped if the corresponding particle list is empty
    for particleList in identifiersExtraInfosDict:
        eventLevelPath = create_path()
        SkipEmptyParticleList = register_module("SkimFilter")
        SkipEmptyParticleList.set_name('SkimFilter_EventLevel_' + particleList)
        SkipEmptyParticleList.param('particleLists', particleList)
        SkipEmptyParticleList.if_true(eventLevelPath, basf2.AfterConditionPath.CONTINUE)
        path.add_module(SkipEmptyParticleList)

        mvaMultipleExperts = register_module('MVAMultipleExperts')
        mvaMultipleExperts.set_name('MVAMultipleExperts_EventLevel_' + particleList)
        mvaMultipleExperts.param('listNames', [particleList])
        mvaMultipleExperts.param('extraInfoNames', [row[0] for row in identifiersExtraInfosDict[particleList]])
        mvaMultipleExperts.param('signalFraction', signalFraction)
        mvaMultipleExperts.param('identifiers', [row[1] for row in identifiersExtraInfosDict[particleList]])
        eventLevelPath.add_module(mvaMultipleExperts)

    if 'KaonPion' in categories and len(identifiersExtraInfosKaonPion) != 0:
        eventLevelKaonPionPath = create_path()
        SkipEmptyParticleList = register_module("SkimFilter")
        SkipEmptyParticleList.set_name('SkimFilter_' + 'K+:inRoe')
        SkipEmptyParticleList.param('particleLists', 'K+:inRoe')
        SkipEmptyParticleList.if_true(eventLevelKaonPionPath, basf2.AfterConditionPath.CONTINUE)
        path.add_module(SkipEmptyParticleList)

        mvaExpertKaonPion = register_module("MVAExpert")
        mvaExpertKaonPion.set_name('MVAExpert_KaonPion_' + 'K+:inRoe')
        mvaExpertKaonPion.param('listNames', ['K+:inRoe'])
        mvaExpertKaonPion.param('extraInfoName', identifiersExtraInfosKaonPion[0][0])
        mvaExpertKaonPion.param('signalFraction', signalFraction)
        mvaExpertKaonPion.param('identifier', identifiersExtraInfosKaonPion[0][1])

        eventLevelKaonPionPath.add_module(mvaExpertKaonPion)

    if mode == 'Sampler':

        for category in categories:
            particleList = AvailableCategories[category].particleList

            methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
            identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'
            targetVariable = 'isRightCategory(' + category + ')'

            if not os.path.isfile(identifierEventLevel):

                if category == 'KaonPion':
                    methodPrefixEventLevelSlowPion = "FlavorTagger_" + getBelleOrBelle2() + \
                                                     "_" + weightFiles + 'EventLevelSlowPionFBDT'
                    identifierEventLevelSlowPion = filesDirectory + '/' + methodPrefixEventLevelSlowPion + '_1.root'
                    if not os.path.isfile(identifierEventLevelSlowPion):
                        B2INFO("Flavor Tagger: event level weight file for the Slow Pion category is absent." +
                               "It is required to sample the training information for the KaonPion category." +
                               "An additional sampling step will be needed after the following training step.")
                        continue

                B2INFO('flavorTagger: file ' + filesDirectory + '/' +
                       methodPrefixEventLevel + "sampled" + fileId + '.root will be saved.')

                ma.applyCuts(particleList, 'isRightCategory(mcAssociated) > 0', path)
                eventLevelpath = create_path()
                SkipEmptyParticleList = register_module("SkimFilter")
                SkipEmptyParticleList.set_name('SkimFilter_EventLevel' + category)
                SkipEmptyParticleList.param('particleLists', particleList)
                SkipEmptyParticleList.if_true(eventLevelpath, basf2.AfterConditionPath.CONTINUE)
                path.add_module(SkipEmptyParticleList)

                ntuple = register_module('VariablesToNtuple')
                ntuple.param('fileName', filesDirectory + '/' + methodPrefixEventLevel + "sampled" + fileId + ".root")
                ntuple.param('treeName', methodPrefixEventLevel + "_tree")
                variablesToBeSaved = getTrainingVariables(category) + [targetVariable, 'ancestorHasWhichFlavor',
                                                                       'isSignal', 'mcPDG', 'mcErrors', 'genMotherPDG',
                                                                       'nMCMatches', 'B0mcErrors']
                if category != 'KaonPion' and category != 'FSC':
                    variablesToBeSaved = variablesToBeSaved + \
                        ['extraInfo(isRightTrack(' + category + '))',
                         'hasHighestProbInCat(' + particleList + ', isRightTrack(' + category + '))']
                ntuple.param('variables', variablesToBeSaved)
                ntuple.param('particleList', particleList)
                eventLevelpath.add_module(ntuple)

    if ReadyMethods != len(categories):
        return False
    else:
        return True

def eventLevelTeacher(weightFiles='B2JpsiKs_mu', categories=None):
    """
    Trains all categories at event level.
    """

    B2INFO('EVENT LEVEL TEACHER')

    ReadyMethods = 0

    if categories is None:
        categories = []

    for category in categories:
        methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
        targetVariable = 'isRightCategory(' + category + ')'
        weightFile = filesDirectory + '/' + methodPrefixEventLevel + "_1.root"

        if os.path.isfile(weightFile):
            ReadyMethods += 1
            continue

        sampledFilesList = glob.glob(filesDirectory + '/' + methodPrefixEventLevel + 'sampled*.root')
        if len(sampledFilesList) == 0:
            B2INFO('flavorTagger: eventLevelTeacher did not find any ' + methodPrefixEventLevel +
                   ".root" + ' file. Please run the flavorTagger in "Sampler" mode afterwards.')

        else:
            B2INFO('flavorTagger: MVA Teacher training' + methodPrefixEventLevel + ' .')
            trainingOptionsEventLevel = basf2_mva.GeneralOptions()
            trainingOptionsEventLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
            trainingOptionsEventLevel.m_treename = methodPrefixEventLevel + "_tree"
            trainingOptionsEventLevel.m_identifier = weightFile
            trainingOptionsEventLevel.m_variables = basf2_mva.vector(*getTrainingVariables(category))
            trainingOptionsEventLevel.m_target_variable = targetVariable
            trainingOptionsEventLevel.m_max_events = maxEventsNumber

            basf2_mva.teacher(trainingOptionsEventLevel, getFastBDTCategories())

            if uploadFlag:
                basf2_mva.upload(weightFile, methodPrefixEventLevel)

    if ReadyMethods != len(categories):
        return False
    else:
        return True

def combinerLevel(mode='Expert', weightFiles='B2JpsiKs_mu', categories=None,
                  variablesCombinerLevel=None, categoriesCombinationCode=None, path=None):
    """
    Samples the input data or tests the combiner according to the selected categories.
    """

    B2INFO('COMBINER LEVEL')

    if categories is None:
        categories = []
    if variablesCombinerLevel is None:
        variablesCombinerLevel = []

    B2INFO("Flavor Tagger: Required Combiner for Categories:")
    for category in categories:
        B2INFO(category)

    B2INFO("Flavor Tagger: which corresponds to a weight file with categories combination code " + categoriesCombinationCode)

    methodPrefixCombinerLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'Combiner' \
        + categoriesCombinationCode

    if mode == 'Sampler':

        if os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root') or \
           os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):
            B2FATAL('flavorTagger: File' + methodPrefixCombinerLevel + 'FBDT' + "_1.root" + ' or ' + methodPrefixCombinerLevel +
                    'FANN' + '_1.root found. Please run the "Expert" mode or delete the file if a new sampling is desired.')

        B2INFO('flavorTagger: Sampling Data on Combiner Level. File' +
               methodPrefixCombinerLevel + ".root" + ' will be saved')

        ntuple = basf2.register_module('VariablesToNtuple')
        ntuple.param('fileName', filesDirectory + '/' + methodPrefixCombinerLevel + "sampled" + fileId + ".root")
        ntuple.param('treeName', methodPrefixCombinerLevel + 'FBDT' + "_tree")
        ntuple.param('variables', variablesCombinerLevel + ['qrCombined'])
        ntuple.param('particleList', "")
        path.add_module(ntuple)

    if mode == 'Expert':

        # Check if weight files are ready
        if TMVAfbdt:
            identifierFBDT = methodPrefixCombinerLevel + 'FBDT'
            if downloadFlag or useOnlyLocalFlag:
                identifierFBDT = filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'

            if downloadFlag:
                if not os.path.isfile(identifierFBDT):
                    basf2_mva.download(methodPrefixCombinerLevel + 'FBDT', identifierFBDT)
                    if not os.path.isfile(identifierFBDT):
                        B2FATAL('Flavor Tagger: Weight file ' + identifierFBDT +
                                ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')

            if useOnlyLocalFlag:
                if not os.path.isfile(identifierFBDT):
                    B2FATAL('flavorTagger: Combinerlevel FastBDT was not trained with this combination of categories.' +
                            ' Weight file ' + identifierFBDT + ' not found. Stopped')

            B2INFO('flavorTagger: Ready to be used with weightFile ' + methodPrefixCombinerLevel + 'FBDT' + '_1.root')

        if FANNmlp:
            identifierFANN = methodPrefixCombinerLevel + 'FANN'
            if downloadFlag or useOnlyLocalFlag:
                identifierFANN = filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'

            if downloadFlag:
                if not os.path.isfile(identifierFANN):
                    basf2_mva.download(methodPrefixCombinerLevel + 'FANN', identifierFANN)
                    if not os.path.isfile(identifierFANN):
                        B2FATAL('Flavor Tagger: Weight file ' + identifierFANN +
                                ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')
            if useOnlyLocalFlag:
                if not os.path.isfile(identifierFANN):
                    B2FATAL('flavorTagger: Combinerlevel FANNMLP was not trained with this combination of categories. ' +
                            ' Weight file ' + identifierFANN + ' not found. Stopped')

            B2INFO('flavorTagger: Ready to be used with weightFile ' + methodPrefixCombinerLevel + 'FANN' + '_1.root')

        # At this stage, all necessary weight files should be ready.
        # Call MVAExpert or MVAMultipleExperts module.
        if TMVAfbdt and not FANNmlp:
            B2INFO('flavorTagger: Apply FBDTMethod ' + methodPrefixCombinerLevel + 'FBDT')
            path.add_module('MVAExpert', listNames=[], extraInfoName='qrCombined' + 'FBDT', signalFraction=signalFraction,
                            identifier=identifierFBDT)

        if FANNmlp and not TMVAfbdt:
            B2INFO('flavorTagger: Apply FANNMethod on combiner level')
            path.add_module('MVAExpert', listNames=[], extraInfoName='qrCombined' + 'FANN', signalFraction=signalFraction,
                            identifier=identifierFANN)

        if FANNmlp and TMVAfbdt:
            B2INFO('flavorTagger: Apply FANNMethod and FBDTMethod  on combiner level')
            mvaMultipleExperts = basf2.register_module('MVAMultipleExperts')
            mvaMultipleExperts.set_name('MVAMultipleExperts_Combiners')
            mvaMultipleExperts.param('listNames', [])
            mvaMultipleExperts.param('extraInfoNames', ['qrCombined' + 'FBDT', 'qrCombined' + 'FANN'])
            mvaMultipleExperts.param('signalFraction', signalFraction)
            mvaMultipleExperts.param('identifiers', [identifierFBDT, identifierFANN])
            path.add_module(mvaMultipleExperts)

def combinerLevelTeacher(weightFiles='B2JpsiKs_mu', variablesCombinerLevel=None,
                         categoriesCombinationCode=None):
    """
    Trains the combiner according to the selected categories.
    """

    B2INFO('COMBINER LEVEL TEACHER')

    if variablesCombinerLevel is None:
        variablesCombinerLevel = []

    methodPrefixCombinerLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'Combiner' \
        + categoriesCombinationCode

    sampledFilesList = glob.glob(filesDirectory + '/' + methodPrefixCombinerLevel + 'sampled*.root')
    if len(sampledFilesList) == 0:
        B2FATAL('FlavorTagger: combinerLevelTeacher did not find any ' +
                methodPrefixCombinerLevel + 'sampled*.root file. Please run the flavorTagger in "Sampler" mode.')

    if TMVAfbdt:

        if not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'):

            B2INFO('flavorTagger: MVA Teacher training a FastBDT on Combiner Level')

            trainingOptionsCombinerLevel = basf2_mva.GeneralOptions()
            trainingOptionsCombinerLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
            trainingOptionsCombinerLevel.m_treename = methodPrefixCombinerLevel + 'FBDT' + "_tree"
            trainingOptionsCombinerLevel.m_identifier = filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + "_1.root"
            trainingOptionsCombinerLevel.m_variables = basf2_mva.vector(*variablesCombinerLevel)
            trainingOptionsCombinerLevel.m_target_variable = 'qrCombined'
            trainingOptionsCombinerLevel.m_max_events = maxEventsNumber

            basf2_mva.teacher(trainingOptionsCombinerLevel, getFastBDTCombiner())

            if uploadFlag:
                basf2_mva.upload(filesDirectory + '/' + methodPrefixCombinerLevel +
                                 'FBDT' + "_1.root", methodPrefixCombinerLevel + 'FBDT')

        elif FANNmlp and not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):

            B2INFO('flavorTagger: Combinerlevel FBDT was already trained with this combination of categories. Weight file ' +
                   methodPrefixCombinerLevel + 'FBDT' + '_1.root has been found.')

        else:
            B2FATAL('flavorTagger: Combinerlevel was already trained with this combination of categories. Weight files ' +
                    methodPrefixCombinerLevel + 'FBDT' + '_1.root and ' +
                    methodPrefixCombinerLevel + 'FANN' + '_1.root has been found. Please use the "Expert" mode')

    if FANNmlp:

        if not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):

            B2INFO('flavorTagger: MVA Teacher training a FANN MLP on Combiner Level')

            trainingOptionsCombinerLevel = basf2_mva.GeneralOptions()
            trainingOptionsCombinerLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
            trainingOptionsCombinerLevel.m_treename = methodPrefixCombinerLevel + 'FBDT' + "_tree"
            trainingOptionsCombinerLevel.m_identifier = filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + "_1.root"
            trainingOptionsCombinerLevel.m_variables = basf2_mva.vector(*variablesCombinerLevel)
            trainingOptionsCombinerLevel.m_target_variable = 'qrCombined'
            trainingOptionsCombinerLevel.m_max_events = maxEventsNumber

            basf2_mva.teacher(trainingOptionsCombinerLevel, getMlpFANNCombiner())

            if uploadFlag:
                basf2_mva.upload(filesDirectory + '/' + methodPrefixCombinerLevel +
                                 'FANN' + "_1.root", methodPrefixCombinerLevel + 'FANN')

        elif TMVAfbdt and not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'):

            B2INFO('flavorTagger: Combinerlevel FBDT was already trained with this combination of categories. Weight file ' +
                   methodPrefixCombinerLevel + 'FANN' + '_1.config has been found.')

        else:
            B2FATAL('flavorTagger: Combinerlevel was already trained with this combination of categories. Weight files ' +
                    methodPrefixCombinerLevel + 'FBDT' + '_1.root and ' +
                    methodPrefixCombinerLevel + 'FANN' + '_1.root has been found. Please use the "Expert" mode')

def getEventLevelParticleLists(categories=None):

    if categories is None:
        categories = []

    eventLevelParticleLists = []

    for category in categories:
        ftCategory = AvailableCategories[category]
        event_tuple = (ftCategory.particleList, ftCategory.eventName, ftCategory.variableName)

        if event_tuple not in eventLevelParticleLists:
            eventLevelParticleLists.append(event_tuple)
        else:
            B2FATAL('Flavor Tagger: ' + category + ' has been already given')

    return eventLevelParticleLists

def flavorTagger(
    particleLists=None,
    mode='Expert',
    weightFiles='B2nunubarBGx1',
    workingDirectory='.',
    combinerMethods=['TMVA-FBDT'],
    categories=[
        'Electron',
        'IntermediateElectron',
        'Muon',
        'IntermediateMuon',
        'KinLepton',
        'IntermediateKinLepton',
        'Kaon',
        'SlowPion',
        'FastHadron',
        'Lambda',
        'FSC',
        'MaximumPstar',
        'KaonPion'],
    maskName='FTDefaultMask',
    saveCategoriesInfo=True,
    useOnlyLocalWeightFiles=False,
    downloadFromDatabaseIfNotFound=False,
    uploadToDatabaseAfterTraining=False,
    samplerFileId='',
    prefix='MC15ri_light-2207-bengal_0',
    useGNN=False,
    identifierGNN='GFlaT_MC15ri_light_2303_iriomote_0',
    path=None,
):
    """
      Defines the whole flavor tagging process for each selected Rest of Event (ROE) built in the steering file.
      The flavor is predicted by Multivariate Methods trained with Variables and MetaVariables which use
      Tracks, ECL- and KLMClusters from the corresponding RestOfEvent dataobject.
      This module can be used to sample the training information, to train and/or to test the flavorTagger.

      @param particleLists                     The ROEs for flavor tagging are selected from the given particle lists.
      @param mode                              The available modes are
                                               ``Expert`` (default), ``Sampler``, and ``Teacher``. In the ``Expert`` mode
                                               Flavor Tagging is applied to the analysis,. In the ``Sampler`` mode you save
                                               save the variables for training. In the ``Teacher`` mode the FlavorTagger is
                                               trained, for this step you do not reconstruct any particle or do any analysis,
                                               you just run the flavorTagger alone.
      @param weightFiles                       Weight files name. Default=
                                               ``B2nunubarBGx1`` (official weight files). If the user self
                                               wants to train the FlavorTagger, the weightfiles name should correspond to the
                                               analysed CP channel in order to avoid confusions. The default name
                                               ``B2nunubarBGx1`` corresponds to
                                               :math:`B^0_{\\rm sig}\\to \\nu \\overline{\\nu}`.
                                               and ``B2JpsiKs_muBGx1`` to
                                               :math:`B^0_{\\rm sig}\\to J/\\psi (\\to \\mu^+ \\mu^-) K_s (\\to \\pi^+ \\pi^-)`.
                                               BGx1 stays for events simulated with background.
      @param workingDirectory                  Path to the directory containing the FlavorTagging/ folder.
      @param combinerMethods                   MVAs for the combiner: ``TMVA-FBDT` (default).
                                               ``FANN-MLP`` is available only with ``prefix=''`` (MC13 weight files).
      @param categories                        Categories used for flavor tagging. By default all are used.
      @param maskName                          Gets ROE particles from a specified ROE mask.
                                               ``FTDefaultMask`` (default): tentative mask definition that will be created
                                               automatically. The definition is as follows:

                                               - Track (pion): thetaInCDCAcceptance and dr<1 and abs(dz)<3
                                               - ECL-cluster (gamma): thetaInCDCAcceptance and clusterNHits>1.5 and \
                                               [[clusterReg==1 and E>0.08] or [clusterReg==2 and E>0.03] or \
                                               [clusterReg==3 and E>0.06]] \
                                               (Same as gamma:pi0eff30_May2020 and gamma:pi0eff40_May2020)

                                               ``all``: all ROE particles are used.
                                               Or one can give any mask name defined before calling this function.
      @param saveCategoriesInfo                Sets to save information of individual categories.
      @param useOnlyLocalWeightFiles           [Expert] Uses only locally saved weight files.
      @param downloadFromDatabaseIfNotFound    [Expert] Weight files are downloaded from
                                               the conditions database if not available in workingDirectory.
      @param uploadToDatabaseAfterTraining     [Expert] For librarians only: uploads weight files to localdb after training.
      @param samplerFileId                     Identifier to paralellize
                                               sampling. Only used in ``Sampler`` mode.  If you are training by yourself and
                                               want to parallelize the sampling, you can run several sampling scripts in
                                               parallel. By changing this parameter you will not overwrite an older sample.
      @param prefix                            Prefix of weight files.
                                               ``MC15ri_light-2207-bengal_0`` (default): Weight files trained for MC15ri samples.
                                               ``''``: Weight files trained for MC13 samples.
      @param useGNN                            Use GNN-based Flavor Tagger in addition with FastBDT-based one.
                                               Please specify the weight file with the option ``identifierGNN``.
                                               [Expert] In the sampler mode, training files for GNN-based Flavor Tagger is produced.
      @param identifierGNN                     The name of weight file of the GNN-based Flavor Tagger.
                                               [Expert] Multiple identifiers can be given with list(str).
      @param path                              Modules are added to this path

    """

    if (not isinstance(particleLists, list)):
        particleLists = [particleLists]  # in case user inputs a particle list as string

    if mode != 'Sampler' and mode != 'Teacher' and mode != 'Expert':
        B2FATAL('flavorTagger: Wrong mode given: The available modes are "Sampler", "Teacher" or "Expert"')

    if len(categories) != len(set(categories)):
        dup = [cat for cat in set(categories) if categories.count(cat) > 1]
        B2WARNING('Flavor Tagger: There are duplicate elements in the given categories list. '
                  << 'The following duplicate elements are removed; ' << ', '.join(dup))
        categories = list(set(categories))

    if len(categories) < 2:
        B2FATAL('Flavor Tagger: Invalid amount of categories. At least two are needed.')
        B2FATAL(
            'Flavor Tagger: Possible categories are  "Electron", "IntermediateElectron", "Muon", "IntermediateMuon", '
            '"KinLepton", "IntermediateKinLepton", "Kaon", "SlowPion", "FastHadron",'
            '"Lambda", "FSC", "MaximumPstar" or "KaonPion" ')

    for category in categories:
        if category not in AvailableCategories:
            B2FATAL('Flavor Tagger: ' + category + ' is not a valid category name given')
            B2FATAL('Flavor Tagger: Available categories are  "Electron", "IntermediateElectron", '
                    '"Muon", "IntermediateMuon", "KinLepton", "IntermediateKinLepton", "Kaon", "SlowPion", "FastHadron", '
                    '"Lambda", "FSC", "MaximumPstar" or "KaonPion" ')

    if mode == 'Expert' and useGNN and identifierGNN == '':
        B2FATAL('Please specify the name of the weight file with ``identifierGNN``')

    # Directory where the weights of the trained Methods are saved
    # workingDirectory = os.environ['BELLE2_LOCAL_DIR'] + '/analysis/data'

    basf2.find_file(workingDirectory)

    global filesDirectory
    filesDirectory = workingDirectory + '/FlavorTagging/TrainedMethods'

    if mode == 'Sampler' or (mode == 'Expert' and downloadFromDatabaseIfNotFound):
        if not basf2.find_file(workingDirectory + '/FlavorTagging', silent=True):
            os.mkdir(workingDirectory + '/FlavorTagging')
            os.mkdir(workingDirectory + '/FlavorTagging/TrainedMethods')
        elif not basf2.find_file(workingDirectory + '/FlavorTagging/TrainedMethods', silent=True):
            os.mkdir(workingDirectory + '/FlavorTagging/TrainedMethods')
        filesDirectory = workingDirectory + '/FlavorTagging/TrainedMethods'

    if len(combinerMethods) < 1 or len(combinerMethods) > 2:
        B2FATAL('flavorTagger: Invalid list of combinerMethods. The available methods are "TMVA-FBDT" and "FANN-MLP"')

    global FANNmlp
    global TMVAfbdt

    FANNmlp = False
    TMVAfbdt = False

    for method in combinerMethods:
        if method == 'TMVA-FBDT':
            TMVAfbdt = True
        elif method == 'FANN-MLP':
            FANNmlp = True
        else:
            B2FATAL('flavorTagger: Invalid list of combinerMethods. The available methods are "TMVA-FBDT" and "FANN-MLP"')

    global fileId
    fileId = samplerFileId

    global useOnlyLocalFlag
    useOnlyLocalFlag = useOnlyLocalWeightFiles

    B2INFO('*** FLAVOR TAGGING ***')
    B2INFO(' ')
    B2INFO('    Working directory is: ' + filesDirectory)
    B2INFO(' ')

    setInteractionWithDatabase(downloadFromDatabaseIfNotFound, uploadToDatabaseAfterTraining)

    if prefix == '':
        set_FlavorTagger_pid_aliases_legacy()
    else:
        set_FlavorTagger_pid_aliases()

    alias_list_for_GNN = []
    if useGNN:
        alias_list_for_GNN = set_GNNFlavorTagger_aliases(categories)

    setInputVariablesWithMask()
    if prefix != '':
        weightFiles = prefix + '_' + weightFiles

    # Create configuration lists and code-name for given category's list
    trackLevelParticleLists = []
    eventLevelParticleLists = []
    variablesCombinerLevel = []
    categoriesCombination = []
    categoriesCombinationCode = 'CatCode'
    for category in categories:
        ftCategory = AvailableCategories[category]

        track_tuple = (ftCategory.particleList, ftCategory.trackName)
        event_tuple = (ftCategory.particleList, ftCategory.eventName, ftCategory.variableName)

        if track_tuple not in trackLevelParticleLists and category != 'MaximumPstar':
            trackLevelParticleLists.append(track_tuple)

        if event_tuple not in eventLevelParticleLists:
            eventLevelParticleLists.append(event_tuple)
            variablesCombinerLevel.append(ftCategory.variableName)
            categoriesCombination.append(ftCategory.code)
        else:
            B2FATAL('Flavor Tagger: ' + category + ' has been already given')

    for code in sorted(categoriesCombination):
        categoriesCombinationCode = categoriesCombinationCode + '%02d' % code

    # Create default ROE-mask
    if maskName == 'FTDefaultMask':
        FTDefaultMask = (
            'FTDefaultMask',
            'thetaInCDCAcceptance and dr<1 and abs(dz)<3',
            'thetaInCDCAcceptance and clusterNHits>1.5 and [[E>0.08 and clusterReg==1] or [E>0.03 and clusterReg==2] or \
                            [E>0.06 and clusterReg==3]]')
        for name in particleLists:
            ma.appendROEMasks(list_name=name, mask_tuples=[FTDefaultMask], path=path)

    # Start ROE-routine
    roe_path = basf2.create_path()
    deadEndPath = basf2.create_path()

    if mode == 'Sampler':
        # Events containing ROE without B-Meson (but not empty) are discarded for training
        ma.signalSideParticleListsFilter(
            particleLists,
            'nROE_Charged(' + maskName + ', 0) > 0 and abs(qrCombined) == 1',
            roe_path,
            deadEndPath)

        FillParticleLists(maskName, categories, roe_path)

        if useGNN:
            if eventLevel('Expert', weightFiles, categories, roe_path):

                ma.rankByHighest('pi+:inRoe', 'p', numBest=0, allowMultiRank=False,
                                 outputVariable='FT_p_rank', overwriteRank=True, path=roe_path)
                ma.fillParticleListFromDummy('vpho:dummy', path=roe_path)
                ma.variablesToNtuple('vpho:dummy',
                                     alias_list_for_GNN,
                                     treename='tree',
                                     filename=f'{filesDirectory}/FlavorTagger_GNN_sampled{fileId}.root',
                                     signalSideParticleList=particleLists[0],
                                     path=roe_path)

        else:
            if eventLevel(mode, weightFiles, categories, roe_path):
                combinerLevel(mode, weightFiles, categories, variablesCombinerLevel, categoriesCombinationCode, roe_path)

        path.for_each('RestOfEvent', 'RestOfEvents', roe_path)

    elif mode == 'Expert':
        # If trigger returns 1 jump into empty path skipping further modules in roe_path
        # run filter with no cut first to get rid of ROEs that are missing the mask of the signal particle
        ma.signalSideParticleListsFilter(particleLists, 'nROE_Charged(' + maskName + ', 0) > 0', roe_path, deadEndPath)

        # Initialization of flavorTaggerInfo dataObject needs to be done in the main path
        flavorTaggerInfoBuilder = basf2.register_module('FlavorTaggerInfoBuilder')
        path.add_module(flavorTaggerInfoBuilder)

        FillParticleLists(maskName, categories, roe_path)

        if eventLevel(mode, weightFiles, categories, roe_path):
            combinerLevel(mode, weightFiles, categories, variablesCombinerLevel, categoriesCombinationCode, roe_path)

            flavorTaggerInfoFiller = basf2.register_module('FlavorTaggerInfoFiller')
            flavorTaggerInfoFiller.param('trackLevelParticleLists', trackLevelParticleLists)
            flavorTaggerInfoFiller.param('eventLevelParticleLists', eventLevelParticleLists)
            flavorTaggerInfoFiller.param('TMVAfbdt', TMVAfbdt)
            flavorTaggerInfoFiller.param('FANNmlp', FANNmlp)
            flavorTaggerInfoFiller.param('qpCategories', saveCategoriesInfo)
            flavorTaggerInfoFiller.param('istrueCategories', saveCategoriesInfo)
            flavorTaggerInfoFiller.param('targetProb', False)
            flavorTaggerInfoFiller.param('trackPointers', False)
            roe_path.add_module(flavorTaggerInfoFiller)  # Add FlavorTag Info filler to roe_path
            add_default_FlavorTagger_aliases()

            if useGNN:
                ma.rankByHighest('pi+:inRoe', 'p', numBest=0, allowMultiRank=False,
                                 outputVariable='FT_p_rank', overwriteRank=True, path=roe_path)
                ma.fillParticleListFromDummy('vpho:dummy', path=roe_path)

                if isinstance(identifierGNN, str):
                    roe_path.add_module('MVAExpert',
                                        listNames='vpho:dummy',
                                        extraInfoName='qrGNN_raw',  # the range of qrGNN_raw is [0,1]
                                        identifier=identifierGNN)

                    ma.variableToSignalSideExtraInfo('vpho:dummy', {'extraInfo(qrGNN_raw)*2-1': 'qrGNN'},
                                                     path=roe_path)

                elif isinstance(identifierGNN, list):
                    identifierGNN = list(set(identifierGNN))

                    extraInfoNames = [f'qrGNN_{i_id}' for i_id in identifierGNN]
                    roe_path.add_module('MVAMultipleExperts',
                                        listNames='vpho:dummy',
                                        extraInfoNames=extraInfoNames,
                                        identifiers=identifierGNN)

                    extraInfoDict = {}
                    for extraInfoName in extraInfoNames:
                        extraInfoDict[f'extraInfo({extraInfoName})*2-1'] = extraInfoName
                        variables.variables.addAlias(extraInfoName, f'extraInfo({extraInfoName})')

                    ma.variableToSignalSideExtraInfo('vpho:dummy', extraInfoDict,
                                                     path=roe_path)

        path.for_each('RestOfEvent', 'RestOfEvents', roe_path)

    elif mode == 'Teacher':
        if eventLevelTeacher(weightFiles, categories):
            combinerLevelTeacher(weightFiles, variablesCombinerLevel, categoriesCombinationCode)

if __name__ == '__main__':

    desc_list = []

    function = globals()["flavorTagger"]
    signature = inspect.formatargspec(*inspect.getfullargspec(function))
    desc_list.append((function.__name__, signature + '\n' + function.__doc__))

    from terminal_utils import Pager
    from basf2.utils import pretty_print_description_list
    with Pager('Flavor Tagger function accepts the following arguments:'):
        pretty_print_description_list(desc_list)


import urllib.request as req
import urllib.error as err
from bs4 import BeautifulSoup
from sys import argv

def main(news_url):

    USAGE_MESSAGE = "="*20 + "\nUSAGE:\npython news_paywall_avoider.py URL\n"+20*"="
    try:
        p_elements = BeautifulSoup(req.urlopen(news_url).read(), 'html.parser').find_all("p")
        for ele in p_elements:
            print(ele.text, end=" ")
    except ValueError as e:
        print(f"\nBAD URL: Might need https:// in front\n{USAGE_MESSAGE}\n{e}")
    except err.HTTPError as e:
        print(f"\nHTTP Error\n{USAGE_MESSAGE}\n{e}")
    except:
        print(f"Unknown Error\n{USAGE_MESSAGE}\n")

if "__main__" == __name__:
    main(argv[1])


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.ops

from ..backbones.basic import ConvUnit
from ..backbones.resnet import ResNet50
from ..backbones.mobilenet import MobileNetV3L
from .components.fpn import FeaturePyramidNetwork
from .operations.anchor import get_priors, make_anchors
from .operations.bbox import clamp_to_canvas, convert_to_cwh, decode_boxes, remove_small, scale_boxes
from .operations.loss import calc_losses, get_losses, match_with_targets
from .operations.post import get_lvidx, final_nms
from .operations.prep import preprocess, prep_targets
from .operations.roi import roi_align_multilevel
from ..utils.weights import load_weights

class RegionProposalNetwork(nn.Module):

    def __init__(self, c, num_anchors, conv_depth):
        super().__init__()
        self.conv = nn.Sequential(*[ConvUnit(c, c, 3, 1, 1, 'relu', None) for _ in range(conv_depth)])
        self.log = nn.Conv2d(c, num_anchors, 1, 1)
        self.reg = nn.Conv2d(c, num_anchors * 4, 1, 1)

    def head(self, x):
        n = x.shape[0]
        x = self.conv(x)
        reg = self.reg(x).permute(0, 2, 3, 1).reshape(n, -1, 4)
        log = self.log(x).permute(0, 2, 3, 1).reshape(n, -1, 1)
        return reg, log

    def filt_dec(self, regs, logs, priors, lvtop):
        res = []
        n = regs[0].shape[0]
        for reg, log, p in zip(regs, logs, priors):
            log, top = log.topk(min(lvtop, log.shape[1]), dim=1)
            reg = reg.gather(1, top.expand(-1, -1, 4))
            pri = p.expand(n, -1, -1).gather(1, top.expand(-1, -1, 4))
            boxes = decode_boxes(reg, pri, mults=(1, 1), clamp=True)
            res.append((boxes, log, log.shape[1]))
        return map(list, zip(*res))

    def forward(self, fmaps, priors, imsizes, cfg, gtboxes=None):
        lvtop = cfg['lvtop'] if not self.training else cfg['lvtop_train']
        imtop = cfg['imtop1'] if not self.training else cfg['imtop1_train']

        tuples = [self.head(x) for x in fmaps]
        regs, logs = map(list, zip(*tuples))

        dregs = [x.detach() for x in regs]
        dlogs = [x.detach() for x in logs]
        boxes, logits, lvlen = self.filt_dec(dregs, dlogs, priors, lvtop)
        boxes = torch.cat(boxes, axis=1)
        obj = torch.cat(logits, axis=1).sigmoid()

        n, dim = boxes.shape[:2]
        boxes, obj = boxes.reshape(-1, 4), obj.flatten()
        idx = torch.nonzero(obj >= cfg['score_thr1']).squeeze()
        boxes, obj = boxes[idx], obj[idx]
        imidx = idx.div(dim, rounding_mode='floor')

        boxes = clamp_to_canvas(boxes, imsizes, imidx)
        boxes, obj, idx, imidx = remove_small(boxes, cfg['min_size1'], obj, idx, imidx)
        groups = imidx * 10 + get_lvidx(idx % dim, lvlen)
        keep = torchvision.ops.batched_nms(boxes, obj, groups, cfg['iou_thr1'])
        keep = torch.cat([keep[imidx[keep] == i][:imtop] for i in range(n)])
        boxes, imidx = boxes[keep], imidx[keep]

        if not self.training:
            return boxes, imidx
        else:
            regs = torch.cat(regs, axis=1)
            logs = torch.cat(logs, axis=1)
            priors = torch.cat(priors)
            loss_obj, loss_reg = get_losses(gtboxes, None, priors, regs, logs, matcher=(0.3, 0.7, True),
                                            sampler=(256, 0.5), types=('ce_bin', 'l1_s'), avg_divs='always_all')
            return boxes, imidx, (loss_obj, loss_reg)

class RoIProcessingNetwork(nn.Module):

    def __init__(self, c, roi_map_size, clin, cfg):
        super().__init__()
        self.conv = nn.ModuleList([ConvUnit(c, c, 3, 1, 1, 'relu') for _ in range(cfg['roi_convdepth'])])
        c1 = c * roi_map_size ** 2
        self.fc = nn.ModuleList([nn.Linear(c1 if i == 0 else clin, clin) for i in range(cfg['roi_mlp_depth'])])
        self.cls = nn.Linear(clin, 1 + cfg['num_classes'])
        self.reg = nn.Linear(clin, cfg['num_classes'] * 4)
        self.ralign_set = cfg['roialign_settings']
        self.bckg_first = cfg['bckg_class_first']

    def heads(self, x):
        if self.conv:
            for layer in self.conv:
                x = layer(x)
        x = x.flatten(start_dim=1)
        if self.fc:
            for mlp in self.fc:
                x = F.relu(mlp(x))
        a = self.reg(x)
        b = self.cls(x)
        return a, b

    def forward(self, proposals, imidx, fmaps, fmaps_strides, imsizes, cfg, gtb=None, gtl=None):
        if self.training:
            # to list
            proposals = [proposals[imidx == i] for i in range(len(gtb))]
            proposals = [torch.cat([p, b.to(torch.float32)]) for p, b in zip(proposals, gtb)]
            # main
            targets, labels, sidx = match_with_targets(gtb, gtl, proposals, 0.5, 0.5, False, 512, 0.25)
            proposals = [p[sampled] for p, sampled in zip(proposals, sidx)]
            # back to joined
            imidx = [torch.full([len(p)], i) for i, p in enumerate(proposals)]
            proposals, imidx = [torch.cat(x) for x in [proposals, imidx]]

        roi_maps = roi_align_multilevel(proposals, imidx, fmaps, fmaps_strides, self.ralign_set)
        reg, log = self.heads(roi_maps)
        reg = reg.reshape(reg.shape[0], -1, 4)

        if self.training:
            labels = torch.cat(labels)
            targets = torch.cat(targets)
            loss_cls = F.cross_entropy(log, labels)
            reg = reg[labels > 0, labels[labels > 0] - 1] # minus 1 because we removed the bckg class
            loss_reg = F.smooth_l1_loss(reg, targets, beta=1/9, reduction='sum') / labels.numel()
            return loss_cls, loss_reg

        scr = F.softmax(log, dim=-1)
        cls = torch.arange(log.shape[1], device=log.device).view(1, -1).expand_as(log)
        scr = scr[:, :-1] if not self.bckg_first else scr[:, 1:]
        cls = cls[:, :-1] if not self.bckg_first else cls[:, 1:]

        n = torch.max(imidx).item() + 1
        dim = reg.shape[1]
        reg, scr, cls = reg.reshape(-1, 4), scr.flatten(), cls.flatten()
        fidx = torch.nonzero(scr > cfg['score_thr2']).squeeze()
        reg, scr, cls = reg[fidx], scr[fidx], cls[fidx]
        idx = fidx.div(dim, rounding_mode='floor')
        proposals, imidx = proposals[idx], imidx[idx]

        proposals = convert_to_cwh(proposals, in_place=True)
        boxes = decode_boxes(reg, proposals, mults=(0.1, 0.2), clamp=True)
        boxes = clamp_to_canvas(boxes, imsizes, imidx)
        boxes, scr, cls, imidx = remove_small(boxes, cfg['min_size2'], scr, cls, imidx)
        b, s, c = final_nms(boxes, scr, cls, imidx, n, cfg['iou_thr2'], cfg['imtop2'])
        return b, s, c

class FasterRCNN(nn.Module):

    thub = 'https://download.pytorch.org/models/'
    mmhub = 'https://download.openmmlab.com/mmdetection/v2.0/'
    links = {
        'tv_resnet50_v1': thub + 'fasterrcnn_resnet50_fpn_coco-258fb6c6.pth',
        'tv_resnet50_v2': thub + 'fasterrcnn_resnet50_fpn_v2_coco-dd69338a.pth',
        'tv_mobilenetv3l_hires': thub + 'fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth',
        'tv_mobilenetv3l_lores': thub + 'fasterrcnn_mobilenet_v3_large_320_fpn-907ea3f9.pth',
        'mm_resnet50': mmhub + 'faster_rcnn/faster_rcnn_r50_fpn_mstrain_3x_coco/'\
                               'faster_rcnn_r50_fpn_mstrain_3x_coco_20210524_110822-e10bd31c.pth',
        'mm_resnet50_animefaces': 'https://github.com/hysts/anime-face-detector/'\
                                  'releases/download/v0.0.1/mmdet_anime-face_faster-rcnn.pth'
    }

    def tv_conversion(self, wd):
        # TorchVision's RoI head predicts boxes for background class too, only to discard
        # them right away, so might as well not calculate it in the first place
        nm = 'roi_heads.box_predictor.bbox_pred.'
        wd[nm + 'weight'] = wd[nm + 'weight'][4:, :] # [364, 1024] -> [360]
        wd[nm + 'bias'] = wd[nm + 'bias'][4:]        # [364] -> [360]
        return wd

    def mm_conversion(self, wd):
        # in MMDet weights for RoI head, representation FC and final reg/log FCs are switched over
        wl = list(wd.items())
        els = [wl.pop(-1) for _ in range(8)][::-1] # last 8 entries
        for el in els[4:] + els[:4]:
            wl.append(el)
        wd = dict(wl)
        return wd

    def config_base(self):
        cfg = {}
        cfg.update(fpn_batchnorm=None, rpn_convdepth=1, roi_convdepth=0, roi_mlp_depth=2)
        cfg.update(resize_min=800, resize_max=1333)
        cfg.update(score_thr1=0.0, iou_thr1=0.7, imtop1=1000, min_size1=0, lvtop=1000)
        cfg.update(score_thr2=0.05, iou_thr2=0.5, imtop2=100, min_size2=0)
        cfg.update(imtop1_train=2000, lvtop_train=2000)
        cfg.update(weights_add_nbatches=False)
        cfg.update(bbone_bn_freeze=False)
        return cfg

    def config_torchvision(self, cfg):
        cfg.update(num_classes=90, bckg_class_first=True, roialign_settings=(2, False))
        cfg.update(weights_sub=None, weights_extra=self.tv_conversion)
        cfg.update(prep_resize='torch', priors_patches='fit', round_anchors=True)
        cfg.update(min_size1=1e-3, min_size2=1e-2)
        return cfg

    def config_mmdet(self, cfg):
        cfg.update(num_classes=80, bckg_class_first=False, roialign_settings=(0, True))
        cfg.update(weights_sub='state_dict', weights_extra=self.mm_conversion)
        cfg.update(prep_resize='cv2', priors_patches='as_is', round_anchors=False)
        return cfg

    def config_resnet(self, cfg):
        cfg.update(bbone='resnet50', resnet_bn_eps=1e-5)
        cfg.update(cins=[256, 512, 1024, 2048], strides=[4, 8, 16, 32, 64])
        cfg.update(anchors=make_anchors([32, 64, 128, 256, 512], [1], [2, 1, 0.5], cfg['round_anchors']))
        return cfg

    def config_mobile(self, cfg):
        cfg.update(bbone='mobilenetv3l')
        cfg.update(cins=[160, 960], strides=[32, 32, 64])
        cfg.update(anchors=make_anchors([32, 32, 32], [1, 2, 4, 8, 16], [2, 1, 0.5], cfg['round_anchors']))
        return cfg

    def get_config(self, src, arch, version=None):
        cfg = self.config_base()
        cfg = self.config_torchvision(cfg) if src == 'tv' else self.config_mmdet(cfg)
        cfg = self.config_resnet(cfg) if arch == 'resnet50' else self.config_mobile(cfg)
        if src == 'tv' and (arch == 'resnet50' and version == 'v1' or arch == 'mobilenetv3l'):
            cfg.update(weights_add_nbatches=True)
            cfg.update(bbone_bn_freeze=True)
        if src == 'tv' and arch == 'resnet50' and version == 'v1':
            cfg.update(resnet_bn_eps=0.0)
        if src == 'tv' and arch == 'resnet50' and version == 'v2':
            cfg.update(fpn_batchnorm=1e-5, rpn_convdepth=2, roi_convdepth=4, roi_mlp_depth=1)
        if src == 'tv' and arch == 'mobilenetv3l' and version == 'lores':
            cfg.update(resize_min=320, resize_max=640, lvtop=150, imtop1=150, score_thr1=0.05)
        if version == 'animefaces':
            cfg.update(num_classes=1)
        return cfg

    def get_backbone(self, cfg):
        if cfg['bbone'] == 'resnet50':
            bn_eps = cfg['resnet_bn_eps']
            bn = bn_eps if not cfg['bbone_bn_freeze'] else (bn_eps, 'frozen')
            return ResNet50(bn=bn, num_freeze=2)
        if cfg['bbone'] == 'mobilenetv3l':
            return MobileNetV3L([13, 16], bn=(1e-05, 'frozen'), num_freeze=7)

    def __init__(self, pretrained='tv_resnet50_v1', device='cpu'):
        super().__init__()
        cfg = self.get_config(*pretrained.split('_'))
        self.cfg = cfg
        self.bases = list(zip(cfg['strides'], cfg['anchors']))
        self.body = self.get_backbone(cfg)
        self.fpn = FeaturePyramidNetwork(cfg['cins'], 256, None, cfg['fpn_batchnorm'], pool=True, smoothP5=True)
        self.rpn = RegionProposalNetwork(256, len(cfg['anchors'][0]), cfg['rpn_convdepth'])
        self.roi = RoIProcessingNetwork(256, 7, 1024, cfg)
        self.to(device)
        load_weights(self, self.links[pretrained], pretrained, device,
                     cfg['weights_extra'], cfg['weights_sub'], cfg['weights_add_nbatches'])

    def forward(self, imgs, targets=None):
        dv = next(self.parameters()).device
        resize = (self.cfg['resize_min'], self.cfg['resize_max'])
        resize_with = self.cfg['prep_resize']
        priors_patches = self.cfg['priors_patches']
        strides = self.cfg['strides']

        x, sz_orig, sz_used = preprocess(imgs, dv, resize, resize_with)
        priors = get_priors(x.shape[2:], self.bases, dv, 'corner', priors_patches, concat=False)
        xs = self.body(x)
        xs = self.fpn(xs)
        if self.training:
            gtb, gtl = prep_targets(targets, sz_used, sz_orig)
            p, imidx, p_losses = self.rpn(xs, priors, sz_used, self.cfg, gtb)
            d_losses = self.roi(p, imidx, xs[:-1], strides[:-1], sz_used, self.cfg, gtb, gtl)
            return (*p_losses, *d_losses)
        else:
            p, imidx = self.rpn(xs, priors, sz_used, self.cfg)
            b, s, c = self.roi(p, imidx, xs[:-1], strides[:-1], sz_used, self.cfg)
            b = scale_boxes(b, sz_orig, sz_used)
            b, s, c = [[t.detach().cpu().numpy() for t in tl] for tl in [b, s, c]]
            return b, s, c


import streamlit as st
from pandas import DataFrame
from common import df_covid_data, df_covid_data_our_countries

def display_datasets_introduction():
    """Display the Importation des données page."""
    st.write(
        "The dataset used for this project is the COVID-19 dataset from \
        Our World in Data: \
        https://github.com/owid/covid-19-data/tree/master/public/data "
    )
    st.write(
        "It is updated daily and includes data on confirmed cases, deaths,\
        hospitalizations, testing, and vaccinations, as well as other \
        variables of potential interest."
    )

def display_dataframe(dataframe: DataFrame):
    """Display a dataframe.

    Parameters
    ----------
    dataframe : DataFrame
        The dataframe to display.
    """
    st.dataframe(dataframe)

def display_datasets_description():
    """Display dataset description."""
    st.write(
        "The dataset with only countries selected for this project,\
        wich are France, Germany, Spain, Italy, Sweden and England:"
    )

def display_datasets_page():
    """Display the page dedicated to the overview of the datasets."""
    st.header("Overview of the datasets")
    display_datasets_introduction()
    display_dataframe(df_covid_data)
    display_datasets_description()
    display_dataframe(df_covid_data_our_countries)

# if __name__ == "__main__":
#     display_datasets_page()


import time
import boto3

ROOT_BUCKET_NAME = "virtual-waiting-room-on-aws"

regions = ['ap-northeast-1', 'ap-south-1', 'ap-southeast-1', 'ca-central-1', \
    'eu-central-1', 'eu-north-1', 'eu-west-1', 'sa-east-1', \
    'us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
print(regions)

for region in regions:
    try:
        s3 = boto3.client('s3', region_name=region)
        name = f"{ROOT_BUCKET_NAME}-{region}"
        response = s3.create_bucket(
            Bucket=name,
            CreateBucketConfiguration={'LocationConstraint': region})
        print(response)
    except Exception as exception:
        print(exception)
    time.sleep(2)


import os
import sys
import datetime
import grading_stats
import re

import discord
from discord.utils import get
from discord.ext import commands
from dotenv import load_dotenv

from sheets_parser import uploadOhTemplate, get_creds, uploadGaTemplate,read_grading_assignments
from grading_stats import getQuestions #,get_message
from reminders import notify_user, notify_user_str

CLASSES = ["CMSC330"]

load_dotenv()
TOKEN = os.getenv('DISCORD_TOKEN')

intents = discord.Intents.all()
intents.members = True

bot = commands.Bot(command_prefix='!', intents=intents)

courseIDs = {}
graders = {}
member_ids = {}
tas = {}
semester = {}

CREDS = get_creds()

def loadCourseIDs():
  for c in CLASSES:
    courseIDs[c] = os.getenv(c)

def loadMemberIDs():
  ids = {}
  for guild in bot.guilds:
    members = guild.members
    for member in members:
      ids[member.name] = member.id
  return ids

def load_graders():
  people = grading_stats.loadGraders()
  for name in people:
    discord_name = people[name][1]
    tas[discord_name] = (people[name][0],name)
    tas[name] = people[name]
  return people

########## USER COMMANDS ##################
'''
@bot.command(name="grading",help='Reponds with what you have left to grade')
async def grader_notify(ctx):
  grader_name = ctx.message.author.name
  if grader_name == "profaccident":
    for g in graders:
      ta = tas[g]
      to_send = single_grader_notify(str(ta[1]),ta)
      if to_send != "":
        await ctx.send(to_send)
  else:
    ta = tas[grader_name]
    to_send = single_grader_notify(grader_name,ta)
    if to_send == "":
      to_send = "Congraulations! You graded everything!"
    await ctx.send(to_send)
'''

########## ADMIN COMMANDS ##################
@bot.command(name="oh",help='creates template of OH schedule')
async def oh_template(ctx):
  uploadOhTemplate(4,CREDS)
  to_send = "Wrote template!"
  await ctx.send(to_send)

@bot.command(name="ga",help='creates template of GA for assignment')
async def ga_template(ctx,assignment):
  match = re.search('\d{3}',str(ctx.guild))
  if not match:
    await ctx.send("You are not in a valid guild or the name needs to change :(")
  else:
    uploadGaTemplate(assignment,CREDS,"CMSC"+match.group())
    to_send = "Wrote " + assignment + " template!"
    await ctx.send(to_send)

@bot.command(name="get_ga",help='gets the grading assignments for assignment')
async def ga_template(ctx,assignment):
  match = re.search('\d{3}',str(ctx.guild))
  if not match:
    await ctx.send("You are not in a valid guild or the name needs to change :(")
  else:
    a = read_grading_assignments(assignment,CREDS,"CMSC"+match.group())
    if a != 0:
      await ctx.send("Error:\n"+a)
    else:
      to_send = "Got grading assignment for " + assignment
      await ctx.send(to_send)

@bot.command(name="notify",help='notifies TAs about grading')
async def ga_template(ctx,assignment):
  match = re.search('\d{3}',str(ctx.guild))
  if not match:
    await ctx.send("You are not in a valid guild or the name needs to change :(")
  else:
    course = "CMSC"+match.group()
    a = get_messages(course,assignment)
    to_send = "Notified TAs about grading " + assignment
    ids=loadMemberIDs()
    #TODO
    await ctx.send(to_send)

################ FUN STUFF ########################33
@bot.event
async def on_message(message):
    if not message.mention_everyone and bot.user.mentioned_in(message):# and message.author.name == 'profaccident':
      emoji = get(message.guild.emojis, name='cringe')
      await message.add_reaction(emoji)
    await bot.process_commands(message)

################ MAIN ########################33

@bot.event
async def on_ready():
  loadCourseIDs()
  global graders
  graders = load_graders()
  global member_ids
  member_ids = loadMemberIDs()
  semester = grading_stats.loadSemester()
  print("done connecting")

bot.run(TOKEN)

############# TRASH #####################
'''
def single_grader_notify(author,ta):
  to_send = notify_grader(ta[0],author)
  if to_send != "":
    to_send = "<@!"+str(member_ids[author])+">"+to_send+"\n"
  return to_send

async def grading_reminder():
  messages = grading_stats.send_notify(semester,graders)
  for course in messages:
    channel = bot.get_channel(int(courseIDs[course]))
    for assignment in messages[course]:
      for name in messages[course][assignment]:
        to_send = notify_grader(ctx,course,name)
        if to_send != "":
          await channel.send("<@!"+name+">\n"+to_send)

def notify_grader(course,name):
  messages = grading_stats.send_notify()
  to_send = ""
  print(messages)
  if course in messages:
    for assignment in messages[course]:
      if name in messages[course][assignment]:
        to_send = to_send + "\nFor "+assignment+":"
        for message in messages[course][assignment][name]:
          to_send = to_send + "\n - "+message
  return to_send
'''


import json
import os
from datetime import datetime
from collections import defaultdict
from flask import Flask, render_template, jsonify, request
from flask_socketio import SocketIO, emit

total_data=defaultdict(list)
prefetch_channellist = list()
i=0
last_used_id=''

app = Flask(__name__)
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
socketio = SocketIO(app)

@app.route("/lastused", methods=["post", "get"])
def lastused():
	global last_used_id
	if request.method == "POST":
		lst= request.form.get("lst");
		last_used_id= lst
	# print(last_used_id)
	return jsonify({"lastused":last_used_id})

@app.route("/channeldata", methods=["post"])
def channeldata():
	title = request.form.get("title")
	if title in total_data:
		return jsonify({"success":True,
			"chatdetail":total_data[title]
			})
	else:
		return jsonify({"success":False
			})

@app.route("/channelname",methods=["GET","POST"])
def channelname():
	print("entered channelname")
	if request.method == "POST":
		title = request.form.get("title")
		prefetch_channellist.append(title)
		print(prefetch_channellist)
	return jsonify({"success":True,
		"prefetch_channellist":prefetch_channellist
		})

@app.route("/")
def index():
    return render_template('index.html')

@app.route("/chats")
def chats():
	# print(total_data)
	return render_template('chat.html')

@socketio.on("send message")
def sendmessage(data):
	global i
	now = datetime.now()
	instance = tuple( (data["username"], data["mes"],
		 now.strftime('%d-%m-%Y'), now.strftime( '%H:%M')) )
	if data["chatroom_name"] not in total_data:
		i=0

	total_data[data["chatroom_name"]].append(instance)
	username = total_data[data["chatroom_name"]][i][0];
	mes = total_data[data["chatroom_name"]][i][1];
	date = total_data[data["chatroom_name"]][i][2];
	time = total_data[data["chatroom_name"]][i][3];
	i += 1
	emit("announce message", {"date":date,
			"mes": mes,
		 	"username":username,
			"time":time},
			 broadcast = True)


import json

from django.http import JsonResponse
from django.shortcuts import render, redirect
from django.views.generic import TemplateView
from django.contrib.auth.mixins import LoginRequiredMixin
from django.contrib import messages
from django.contrib.auth.decorators import login_required

#Para conectarse a mongoDB desde python
from utils import *
from bson import ObjectId
import pymongo
#Cast from str to array
from ast import literal_eval

# Create your views here.

def delete_course(request, id):
    db = connect("proesCol")

    query = {"_id": ObjectId(id)}
    print(query)
    db.delete_one(query)
    return redirect('courses')

@login_required
def Courses(request):
    # Para saber las asignaturas de un profesor
    db = connect("proesCol")
    if request.method == 'POST':
        #"""""
        my_id = request.POST["id_materia"]
        nuevo_nombre = request.POST["nuevo_nombre"]
        query = {"_id": ObjectId(my_id)}

        new_values = {"$set":{"asignatura.nombre":nuevo_nombre}}
        db.update_one(query, new_values)
        #"""
        return redirect('courses')
    profesor = request.user.username
    asignaturas = []
    salida = db.find({"profesor": profesor}, {"asignatura.nombre": 1})
    for x in salida:
        x["id"] = str(x["_id"])
        asignaturas.append(x)
    return render(request, "courses/courses.html", {"contexto": asignaturas})  # "about.html"

def TopWorst(request, id):
    #se debe recibir un id en el argumento, cuando esté listo se debe borrar el id de abajo
    #Solo funciona por ahora con una sola materia
    # id = '6201dd30cf634ef49138dec8'
    db = connect("proesCol")
    salida = db.find_one({"_id":ObjectId(id)},{"estudiantes":1,"_id":0})["estudiantes"]

    ordered = sorted(salida, key=lambda d: d['definitiva'])
    salidaArr=[]

    if len(ordered)>= 3:
        salidaArr={"est1":ordered[0],"est2":ordered[1],"est3":ordered[2]}
    else:
        for i in range(len(ordered)):
            salidaArr.append(ordered[i])
    return render(request,"courses/topworst.html",{"contexto":salidaArr})  # "about.html"

def TopBest(request, id):
    #se debe recibir un id en el argumento, cuando esté listo se debe borrar el id de abajo
    #Solo funciona por ahora con una sola materia
    # id = '6201dd30cf634ef49138dec8'
    db = connect("proesCol")
    salida = db.find_one({"_id":ObjectId(id)},{"estudiantes":1,"_id":0})["estudiantes"]

    ordered = sorted(salida, key=lambda d: d['definitiva'], reverse=True)
    salidaArr=[]

    if len(ordered)>= 3:
        salidaArr={"est1":ordered[0],"est2":ordered[1],"est3":ordered[2]}
    else:
        for i in range(len(ordered)):
            salidaArr.append(ordered[i])
    return render(request,"courses/topbest.html",{"contexto":salidaArr})  # "about.html"

@login_required
def NewCourse(request):
    #id = str(db.find_one({"asignatura.nombre":"ESTADÍSTICA"},{"_id":1})["_id"])
    #out = db.find_one({"_id":ObjectId(id)})
    if request.method == 'POST':
        nombre = request.POST["nombre"]
        min = float(request.POST["min"])
        max = float(request.POST["max"])
        profesor = request.user.username
        query = {"asignatura":{"nombre":nombre,"min":min,"max":max},"profesor":profesor,"tipo_notas":[],"estudiantes":{}}
        db = connect("proesCol")
        db.insert(query)
        return redirect('courses')
    return render(request, "courses/newcourse.html")  # "about.html"

@login_required
def Spreadsheet(request, id):
    #tipo_notas y estudiantes es lo que se va a recibir del spreedsheet

    db = connect("proesCol")
    query = {"_id":ObjectId(id)}
    salida= db.find(query)[0]

    if request.method == 'POST':
        estudiantes = literal_eval(request.POST['estudiantes'])
        tipo_notas = literal_eval(request.POST['tipo_notas'])
        if tipo_notas != salida["tipo_notas"]:
            db.update_one(query,{"$set":{"tipo_notas":tipo_notas}})
        if estudiantes != salida["estudiantes"]:
            db.update_one(query,{"$set":{"estudiantes":estudiantes}})
        salida= db.find(query)[0]

    salida["id"] = str(salida["_id"])
    return render(request, "courses/spreadsheet.html",{"contexto":salida})   # "about.html"


"""
https://www.acmicpc.net/problem/17472

7 8
0 0 0 0 0 0 1 1
1 1 0 0 0 0 1 1
1 1 0 0 0 0 0 0
1 1 0 0 0 1 1 0
0 0 0 0 0 1 1 0
0 0 0 0 0 0 0 0
1 1 1 1 1 1 1 1
"""

import collections

N, M = map(int, input().split())
matrix = [list(map(int, input().split())) for _ in range(N)]

dr = (0, 1, 0, -1)
dc = (1, 0, -1, 0)

def bfs(r: int, c: int, visited: set):
    deque = collections.deque([(r, c)])
    island = set([(r, c)])

    while deque:
        _r, _c = deque.popleft()

        for i in range(4):
            nr = _r + dr[i]
            nc = _c + dc[i]
            if (
                (0 <= nr < N)
                and (0 <= nc < M)
                and (nr, nc) not in visited
                and matrix[nr][nc] == 1
            ):
                visited.add((nr, nc))
                deque.append((nr, nc))
                island.add((nr, nc))

    return list(island)

def find_islands():
    """
    Find all islands using BFS
    """
    visited = set()
    islands = []

    for r in range(N):
        for c in range(M):
            if matrix[r][c] == 1 and (r, c) not in visited:
                islands.append(bfs(r, c, visited))

    return islands

def build_bridges(id: int, island: list):
    def can_go(r, c):
        return (0 <= r < N) and (0 <= c < M) and matrix[r][c] != 1

    def find_reached_island(r, c):
        for i, island in enumerate(islands):
            if (r, c) in island:
                return i

        return -1

    bridges = set()
    for r, c in island:
        for i in range(4):
            nr = r + dr[i]
            nc = c + dc[i]

            len = 0
            while can_go(nr, nc):
                len += 1

                nr = nr + dr[i]
                nc = nc + dc[i]

            reached_island = find_reached_island(nr, nc)

            if len < 2 or reached_island == -1:
                continue

            bridges.add((id, reached_island, len))

    return bridges

def find_parent(x):
    if parent[x] != x:
        parent[x] = find_parent(parent[x])
    return parent[x]

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a != b:
        parent[a] = b

def kruskal():
    answer, count = 0, 0

    for a, b, length in bridges:
        if find_parent(a) != find_parent(b):
            union_parent(a, b)

            answer += length
            count += 1

    return answer, count

"""
Solve Problem
"""
islands = find_islands()
bridges_dict = collections.defaultdict(list)
bridges = []
for id, island in enumerate(islands):
    _bridges = build_bridges(id, island)
    for _bridge in _bridges:
        a, b, length = _bridge
        bridges_dict[a].append((b, length))
        bridges.append(_bridge)

parent = [i for i in range(len(islands) + 1)]
bridges.sort(key=lambda x: x[2])

answer, count = kruskal()

print(answer if count == len(islands) - 1 else -1)


import requests
ua = {"User-Agent":"Mozilla/5.0"}
import os
url = "http://pic70.nipic.com/file/20150626/610812_114205201599_2.jpg"
root = "D://PythonWeb//"
path = root + url.split('/')[-1]

try:
    if not os.path.exists(root):
        os.mkdir(root)
    if not os.path.exists(path):
        r = requests.get(url, headers = ua)
        with open(path, 'wb') as f:
            f.write(r.content)
            f.close()
            print("文件保存成功")
    else:
        print("文件已存在")
except:
    print ("爬取失败")


import os
import re
import syslog
import tempfile
import time
import fcntl
from configparser import ConfigParser
import requests

def uri_reader(uri):
    req_opts = {
        'url': uri,
        'timeout': 120,
        'stream': True
    }
    try:
        req = requests.get(**req_opts)
    except Exception as e:
        syslog.syslog(syslog.LOG_ERR,'blacklist download : unable to download file from %s (error : %s)' % (uri, e))
        return

    if req.status_code >= 200 and req.status_code <= 299:
        req.raw.decode_content = True
        prev_chop  = ''
        while True:
            chop = req.raw.read(1024).decode()
            if not chop:
                if prev_chop:
                    yield prev_chop
                break
            else:
                parts = (prev_chop + chop).split('\n')
                if parts[-1] != "\n":
                    prev_chop = parts.pop()
                else:
                    prev_chop = ''
                for part in parts:
                    yield part
    else:
        syslog.syslog(syslog.LOG_ERR,
            'blacklist download : unable to download file from %s (status_code: %d)' % (uri, req.status_code)
        )

if __name__ == '__main__':
    # check for a running download process, this may take a while so it's better to check...
    try:
        lck = open('/tmp/unbound-download_blacklists.tmp', 'w+')
        fcntl.flock(lck, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        # already running, exit status 99
        sys.exit(99)

    startup_time = time.time()
    syslog.openlog('unbound', logoption=syslog.LOG_DAEMON, facility=syslog.LOG_LOCAL4)
    blacklist_items = set()
    if os.path.exists('/var/unbound/etc/blacklists.ini'):
        cnf = ConfigParser()
        cnf.read('/var/unbound/etc/blacklists.ini')
        # exclude (white) lists, compile to regex to be used to filter blacklist entries
        if cnf.has_section('exclude'):
            exclude_list = set()
            for exclude_item in cnf['exclude']:
                try:
                    re.compile(cnf['exclude'][exclude_item], re.IGNORECASE)
                    exclude_list.add(cnf['exclude'][exclude_item])
                except re.error:
                    syslog.syslog(syslog.LOG_ERR,
                        'blacklist download : skip invalid whitelist exclude pattern "%s" (%s)' % (
                            exclude_item, cnf['exclude'][exclude_item]
                        )
                    )
            if not exclude_list:
                exclude_list.add('$^')

            wp = '|'.join(exclude_list)
            whitelist_pattern = re.compile(wp, re.IGNORECASE)
            syslog.syslog(syslog.LOG_NOTICE, 'blacklist download : exclude domains matching %s' % wp)

        # fetch all blacklists
        if cnf.has_section('blacklists'):
            for blacklist in cnf['blacklists']:
                file_stats = {'uri': cnf['blacklists'][blacklist], 'skip' : 0, 'blacklist': 0, 'lines' :0}
                for line in uri_reader(cnf['blacklists'][blacklist]):
                    file_stats['lines'] += 1
                    # cut line into parts before comment marker (if any)
                    tmp = line.split('#')[0].split()
                    entry = None
                    while tmp:
                        entry = tmp.pop(-1)
                        if entry not in ['127.0.0.1', '0.0.0.0']:
                            break
                    if entry:
                        if whitelist_pattern.match(entry):
                            file_stats['skip'] += 1
                        else:
                            file_stats['blacklist'] += 1
                            blacklist_items.add(entry)

                syslog.syslog(
                    syslog.LOG_NOTICE,
                    'blacklist download %(uri)s (lines: %(lines)d exclude: %(skip)d black: %(blacklist)d' % file_stats
                )

    # write out results
    with open("/var/unbound/etc/dnsbl.conf", 'w') as unbound_outf:
        if blacklist_items:
            unbound_outf.write('server:\n')
            for entry in blacklist_items:
                unbound_outf.write("local-data: \"%s A 0.0.0.0\"\n" % entry)

    syslog.syslog(syslog.LOG_NOTICE, "blacklist download done in %0.2f seconds (%d records)" % (
        time.time() - startup_time, len(blacklist_items)
    ))