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)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This module is intended as a quick and dirty drop in replacement of kombu module, covering only
the functionality we use in zmon.
"""

import re
from collections import namedtuple

Connection = namedtuple('Connection', 'hostname port virtual_host')

def parse_redis_conn(conn_str):
    '''
    Emulates kombu.connection.Connection that we were using only to parse redis connection string
    :param conn_str: example 'redis://localhost:6379/0'
    :return: namedtuple(hostname, port, virtual_host)

    >>> parse_redis_conn('localhost:6379')
    Connection(hostname='localhost', port=6379, virtual_host='0')

    >>> parse_redis_conn('localhost:6379/0')
    Connection(hostname='localhost', port=6379, virtual_host='0')
    '''

    if '://' not in conn_str:
        conn_str = 'redis://' + conn_str

    conn_regex = r'redis://([-.a-zA-Z0-9_]+):([0-9]+)(/[0-9]+)?'
    m = re.match(conn_regex, conn_str)
    if not m:
        raise Exception('unable to parse redis connection string: {}'.format(conn_str))
    return Connection(m.group(1), int(m.group(2)), m.group(3).lstrip('/') if m.group(3) else '0')


import cv2
from datetime import datetime
from get_db import get_user_profiles
import face_recognition
import datetime as dt
from collections import Counter
import itertools
from Insert_Measure_Info import Insert_Measure_Info
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from pathlib import Path
from Combine_database import fetch_newest_temperature_db, Update_Measure_Info
from get_file_path import get_file_path

# Return:
# min_distance: smallest distance
# id: the id of the min_distance

def find_closest(list_of_face_encodings, unknown_face_encoding):
    distances = face_recognition.face_distance(
        list_of_face_encodings, unknown_face_encoding)
    min_distance = min(distances)
    id = distances.argmin()
    return id, min_distance

# Get the indices of faces in a list of frames which
#  corresponding to each face in the last frame.
def same_face_indices(prev_encodings, prev_locations, tolerance):
    selected_bools = []
    for col, x in enumerate(prev_encodings):
        selected_bools.append([])
        for y in x:
            selected_bools[col].append(False)
    same_indices = [[] for x in range(len(prev_encodings[-1]))]
    for i, (last_encoding, last_location, indices) \
            in enumerate(zip(
                prev_encodings[-1],
                prev_locations[-1],
                same_indices
            )):
        for row, (encodings, locations, bools)\
                in enumerate(zip(
                    prev_encodings[:-1],
                    prev_locations[:-1],
                    selected_bools[:-1]
                )):
            closest = None
            min_distance = 1.0
            for col in range(len(encodings)):
                distance = face_recognition.face_distance(
                    [encodings[col]], last_encoding
                )
                if distance < min_distance and distance < tolerance\
                        and bools[col] is False:
                    closest = col
                    min_distance = distance
            if closest is not None:
                bools[closest] = True
                indices.append((row, closest))
        indices.append((len(prev_encodings) - 1, i))
    return same_indices

def draw_results(frame, locations, results, user_names,
                 tolerance, font, scale=2):
    rgb_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    for (top, right, bottom, left), (_, min_distance), user_name \
            in zip(locations, results, user_names):
        top *= scale
        right *= scale
        bottom *= scale
        left *= scale

        draw = ImageDraw.Draw(rgb_frame)
        rectColor = (255, 0, 0)
        draw.rectangle([(left, top), (right, bottom)],
                       fill=None, outline=rectColor, width=5)

        draw.rectangle([(left, bottom - 70), (right, bottom)],
                       fill=rectColor, outline=None)
        textColor = (255, 255, 255)
        draw.text((left + 6, bottom - 72),
                  str(f'{user_name}, Diff: {min_distance: > .4f}'),
                  font=font, fill=textColor)
        match_state = "Not Matched" if min_distance > tolerance else "Matched"
        draw.text((left + 6, bottom - 36),
                  str(f' {match_state}'),
                  font=font, fill=textColor)
    result_frame = cv2.cvtColor(np.asarray(rgb_frame), cv2.COLOR_RGB2BGR)
    return result_frame

def main():
    database_name = './Release/teacher.db'
    tolerance = 0.35
    frame_buffer_size = 10  # number of buffered frames to generate result
    # [[en_face, ID, Name], ...]
    user_profile_list = get_user_profiles(database_name)
    known_face_encodings = [x[0] for x in user_profile_list]
    guest_encodings = []

    font = ImageFont.truetype(
        str(Path(get_file_path()).joinpath('NotoSansCJK-Regular222.ttc')), 20)
    prev_locations = []  # Previous face locations in buffered frames
    prev_encodings = []  # Previous face encodings in buffered frames
    # Indices of previous matched encodings in buffered frames
    prev_matched_ids = []
    time_dict = {}  # {id: leaving_time}
    last_detected_time = dt.datetime.min
    results = []  # [[closest_id, distance], ...]

    guest_count = 0
    user_names = []  # name of detected people in the current frame
    user_name_dict = {}  # {id: name}
    record_frame_count = 0

    video_capture = cv2.VideoCapture(1)
    for frame_count in itertools.count():
        _, frame = video_capture.read()

        # For every two frames, Skip one frame.
        if frame_count % 2 == 0:
            small_frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
            rgb_small_frame = small_frame[:, :, ::-1]
            new_locations = face_recognition.face_locations(rgb_small_frame)
            if len(new_locations) > 0 or (
                dt.datetime.now() - last_detected_time
            ).total_seconds() > 2.0:
                locations = new_locations
                last_detected_time = dt.datetime.now()
                encodings = face_recognition.face_encodings(
                    rgb_small_frame, locations)
                matched_ids = []
                for encoding in encodings:
                    id, distance = find_closest(known_face_encodings, encoding)
                    if distance > tolerance and len(guest_encodings) > 0:
                        id2, distance2 = find_closest(
                            guest_encodings, encoding)
                        id2 += len(known_face_encodings)
                        if distance2 < distance:
                            id, distance = id2, distance2
                    matched_ids.append(id)
                start = 0 if len(prev_locations) < frame_buffer_size else 1
                prev_locations[0:] = prev_locations[start:] + [locations]
                prev_encodings[0:] = prev_encodings[start:] + [encodings]
                prev_matched_ids[0:] = prev_matched_ids[start:] + [matched_ids]
                record_frame_count += 1
                # Generate results every frame_buffer_size frames
                if (record_frame_count % frame_buffer_size
                        == frame_buffer_size - 1):
                    # Get the indices of encodings that have the most matching
                    indices_list = same_face_indices(
                        prev_encodings, prev_locations, tolerance)
                    results = []
                    for encoding, location, indices in zip(
                            encodings, locations, indices_list):
                        ids = [prev_matched_ids[row][col]
                               for row, col in indices]
                        id = Counter(ids).most_common(1)[0][0]
                        distance = face_recognition.face_distance(
                            (known_face_encodings + guest_encodings)[id: id+1],
                            encoding)[0]
                        results.append(
                            (id, distance)
                        )
                    user_names = []
                    now = datetime.now()
                    for i, (id, min_distance) in enumerate(results):
                        # Is new guest
                        if min_distance > tolerance:
                            user_name = f'訪客{guest_count + 1}'
                            id = len(known_face_encodings) + \
                                len(guest_encodings)
                            guest_encodings.append(encodings[i])
                            guest_count += 1
                        # Is in database
                        elif id < len(known_face_encodings):
                            user_name = user_profile_list[id][2]
                        # Is old guest
                        else:
                            user_name = user_name_dict[id]

                        # 若現在距離人臉最後偵測時間大於十秒，將資料寫進資料庫

                        if (now - time_dict.setdefault(
                            user_name, dt.datetime.min))\
                                .total_seconds() > 10.0:
                            Insert_Measure_Info(
                                database_name, [user_name, now])
                            data = fetch_newest_temperature_db(database_name)
                            Update_Measure_Info(database_name, data)
                            print('寫入')

                        print(user_name)
                        time_dict[user_name] = now
                        user_name_dict[id] = user_name
                        user_names.append(user_name)
        # TODO: Match locations with results
        result_frame = draw_results(
            frame, locations, results, user_names, tolerance, font)
        cv2.imshow('Video', result_frame)
        # Press q to quit
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

if __name__ == '__main__':
    main()


from tqdm import tqdm

from MCTS import MCTS
from BinaryTree import BinaryTree
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(15)

def run_experiment(max_iterations, dynamic_c=False):
    """
    Run a single experiment of a sequence of MCTS searches to find the optimal path.
    :param max_iterations: Number of iterations to run the MCTS.
    :param dynamic_c: Boolean indicating whether to use a dynamic c or not.
    :return: value of the optimal path found from the search
    """
    tree = BinaryTree(depth=12, b=20, tau=3)
    best_leaf = max(tree.leaves)
    mcts = MCTS(max_iterations=max_iterations, c=2)

    optimal_path = ""
    while tree.depth > 0:
        # search the best direction
        direction = mcts.search(tree, dynamic_c=dynamic_c, verbose=False)
        optimal_path += direction

        # update the tree
        tree.update_root(direction)

    # return the distance of the optimal path found from the search wrt the best leaf
    return sum(1 for a, b in zip(optimal_path, best_leaf.address) if a != b)

def main():
    # compute statistics for static c and dynamic c
    n_iterations = np.logspace(0.7, 3, num=18, base=10, dtype=int)
    values_static_c = [run_experiment(max_iterations=n, dynamic_c=False) for n in tqdm(n_iterations, desc='Execute MCTS with c=2', unit=' experiment')]
    values_dynamic_c = [run_experiment(max_iterations=n, dynamic_c=True) for n in tqdm(n_iterations, desc='Execute MCTS with dynamic c', unit=' experiment')]

    # plot the results
    plt.figure(figsize=(8, 4))
    plt.plot(n_iterations, values_dynamic_c, '-o', label="MCTS with dynamic c")
    plt.plot(n_iterations, values_static_c, '-o', label="MCTS with c=2")
    plt.xlabel("Number of iterations")
    plt.ylabel("Distance of the optimal path from the best leaf")
    plt.title("Compare the value of the optimal path found by MCTS with and without dynamic c")
    plt.grid(linestyle='--', linewidth=1)
    plt.xscale("log")
    plt.xticks(n_iterations, n_iterations)
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()


#!/usr/bin/env python

"""
Author:  Giovanni merlos Mellini
License: GNU General Public License v3.0
         https://github.com/gmellini/twitter-scraper/blob/master/LICENSE

This script allows you to get a complete list of twitter threads replies
so you can have a fast and complete view of complex threads even if you
are not cited in all the tweet branches

Check https://github.com/gmellini/twitter-scraper to read more on this

= Twitter Access Tokens =
Generate here
 https://developer.twitter.com/en/docs/basics/authentication/guides/access-tokens.html

and modify [ CONSUMER_KEY, CONSUMER_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET ]

= LIMITATION =
Twitter search API only returns results from last 7 days. This means that
search results are limited to last 7 days

= TWITTER-SCRAPER-PY =
twitter-scraper.py reads a text file fileed with tweets URL (one per line) in
the following format https://twitter.com/<SCREEN_NAME>/status/<ID>
The script checks and writes to stdout any reply to the given tweet, keeping
replies indentation.

Option -s gives a short output that can be useful to diff content between
different iterations of the script; this way you can check for newer replies
and notify.

usage: twitter-scraper.py -f file [-s]
Options:
-f  : name of the input file that contains twitter URLs (1 per line) in the following
      format: https://twitter.com/<SCREEN_NAME>/status/<ID>
-s  : csv output; useful to diff content between different iterations of the script

= REQUIREMENTS =
The script is tested with python 2.7 and 3.6 on Ubuntu 18.04 and 18.10
- install required pip packages:
    # Python 2.7
    sudo pip install python-twitter
    sudo pip install pytz
    # Python 3
    sudo pip3 install python-twitter
    sudo pip3 install pytz
- modify your timezone to have local dates
    - Full list here: https://gist.github.com/heyalexej/8bf688fd67d7199be4a1682b3eec7568
    - Current: local_timezone = 'Europe/Rome'

= EXAMPLES =
$ ./twitter-scraper.py  -f tweet.list
[...]
$ ./twitter-scraper.py  -f tweet.list -s
[...]

= CREDITS =
Based on the initial work made by @edsu
 https://gist.github.com/edsu/54e6f7d63df3866a87a15aed17b51eaf
"""

import sys
import getopt
import json
import time
import twitter
import re
import os.path
import locale
from os import environ as e
from datetime import datetime
from pytz import timezone
try:
    import urllib # python 2.7
except ImportError:
    import urllib.parse # python 3

################################################
# SET YOUR TWITTER KEYS/SECRETS
e["consumer_key"]="CONSUMER_KEY"
e['consumer_secret']="CONSUMER_SECRET"
e["access_token_key"]="ACCESS_TOKEN"
e['access_token_secret']="ACCESS_TOKEN_SECRET"
# SET TIMEZONE FOR OUTPUT
local_timezone = 'Europe/Rome'
################################################

# Global vars
t = twitter.Api(
    consumer_key=e["consumer_key"],
    consumer_secret=e["consumer_secret"],
    access_token_key=e["access_token_key"],
    access_token_secret=e["access_token_secret"],
    sleep_on_rate_limit=True
)

tweet_url = None
tweet_user = None
twitter_id = None
twitter_sceen_name = None
short_output = False

# True to have debug messages
debug = False

def print_usage():
    print ('usage: twitter-scraper.py -f file [-s]')
    print ('Options:')
    print ('-f\t: name of the input file that contains twitter URLs (1 per line) in the following format: https://twitter.com/<SCREEN_NAME>/status/<ID>')
    print ('-s\t: csv output; useful to diff content between different iterations of the script')

def get_tweets(filename):
    global twitter_id
    global twitter_sceen_name
    global tweet_url
    global short_output

    for line in open(filename):
        if not short_output: print('=========================================')
        line.strip()
        if not re.search('^#', line):
            m = re.search('https://twitter.com/(.*)/status/(.*)', line)
            if m:
                twitter_sceen_name = m.group(1)
                twitter_id = m.group(2)
            else:
        	    if not short_output: print('[WARNING] Malformed URL - Not scraping %s' % (line))
        	    continue
            twitter_json = '{"user":{"screen_name": "'+twitter_sceen_name+'"},"id": '+twitter_id+'}'
            tweet_url = 'https://twitter.com/%s/status/%s' % (twitter_sceen_name, twitter_id)
            if not short_output:
                print('[INFO] Start scraping from tweet URL %s' % (tweet_url))
                print('')
            yield twitter.Status.NewFromJsonDict(json.loads(twitter_json))
        else:
            if not short_output: print('[INFO] Comment - %s' % (line))
            continue

def get_replies(tweet):
    global tweet_user
    global debug
    global short_output

    tweet_user = tweet.user.screen_name
    tweet_id = tweet.id
    max_id = None
    page_index = 0
    while True:
        page_index += 1
        try:
            term = "to:%s" % tweet_user
            if debug: print('[DEBUG] Search: term=%s since_id=%s max_id=%s count=100' % (term, tweet_id, max_id))
            replies = t.GetSearch(
                term=term,
                since_id=tweet_id,
                max_id=max_id,
                count=100,
            )
        except twitter.error.TwitterError as e:
            if not short_output: print('[WARNING] Caught twitter api error, sleep for 60s: %s', (e))
            time.sleep(60)
            continue

        if not replies:
            if debug: print('[DEBUG] Found no replies. Search continues...')
            if debug: print('-----------------------------------------')
            break

        for reply in replies:
            if reply.in_reply_to_status_id == tweet_id:
                if debug: print('[DEBUG] Reply found')
                yield reply
                # recursive magic to also get replies to the reply
                for reply_to_reply in get_replies(reply):
                    yield reply_to_reply
            max_id = reply.id

        max_id = reply.id - 1

        if len(replies) != 100:
            break

def main(argv):
    global tweet_url
    global twitter_id
    global short_output
    global local_timezone

    # Parse ARGV[]
    try:
        opts, args = getopt.getopt(argv, "hsf:")
    except getopt.GetoptError:
        print_usage()
        sys.exit(2)

    find_f = False
    for opt, arg in opts:
        if opt == '-h':
            print_usage()
            sys.exit(0)
        elif opt == '-f':
            file_in = arg
            find_f = True
        elif opt == '-s':
            short_output = True

    if find_f is False:
        print('[ERROR] Missing mandatory option -f')
        print('')
        print_usage()
        sys.exit(2)

    if not os.path.exists(file_in):
        print('[ERROR] File %s not exists' % (file_in))
        print('')
        print_usage()
        sys.exit(2)

    # Verify key/secrets
    if e["consumer_key"] == "CONSUMER_KEY" or e['consumer_secret'] == "CONSUMER_SECRET" or e["access_token_key"] == "ACCESS_TOKEN" or e['access_token_secret'] == "ACCESS_TOKEN_SECRET":
        print('[ERROR] It seems that Twitter key/secrets where not set.')
        print('\tGenerate here https://developer.twitter.com/en/docs/basics/authentication/guides/access-tokens.html and set to start playing')
        print('')
        print_usage()
        sys.exit(2)

    if short_output: print('date,reply,parent_thread');

    for tweet in get_tweets(file_in):
        if tweet_url is not None:
            tabs = ''
            last = None
            count = 0
            for reply in get_replies(tweet):
                j = json.loads(reply.AsJsonString())

                if int(reply.in_reply_to_status_id) == int(twitter_id):
                	tabs = ''
                elif int(reply.in_reply_to_status_id) == int(last):
                    tabs = '%s  ' % tabs
                else:
                    tabs = '%s\b\b' % tabs

                last = j['id']
                pre_human_date = j['created_at'].replace('+0000 ', '')
                human_date = datetime.strptime(pre_human_date, '%a %b %d %H:%M:%S %Y').replace(tzinfo=timezone('UTC')).astimezone(tz=timezone(local_timezone)).strftime('%d/%m/%Y %H:%M:%S')
                count += 1

                if not short_output:
                    print('%s/----------------------------------' % (tabs))
                    print('%s| From:\t %s (@%s)' % (tabs, j['user']['name'], j['user']['screen_name']))
                    if debug: print('%s| Twitter date:\t %s' % (tabs, j['created_at']))
                    print('%s| Date:\t %s' % (tabs, human_date))
                    print('%s| URL:\t https://twitter.com/%s/status/%s' % (tabs, j['user']['screen_name'], j['id']))
                    print('%s| %s' % (tabs, j['text']))
                    print('%s\\----------------------------------' % (tabs))
                    print ('')

                if short_output: print('%s,https://twitter.com/%s/status/%s,%s' % (human_date, j['user']['screen_name'], j['id'], tweet_url))

                if debug:
                    print('')
                    print('[DEBUG] status ID: %s' % (j['id']))
                    print('[DEBUG] in reply to: %s' % (j['in_reply_to_status_id']))
                    print('[DEBUG] JSON response:')
                    print(reply.AsJsonString())
            if not short_output:
                print('')
                print('[INFO] Total replies: %s' % (count))

    if not short_output: print('=========================================')

if __name__ == "__main__":
    main(sys.argv[1:])
    sys.exit(0)


import os
import abc

import numpy as np
import scipy.linalg
import scipy.signal as sig

# TODO: Clean up this interface
class Basis(object):
    __metaclass__ = abc.ABCMeta

    def __init__(self, B, dt, dt_max,
                 orth=False,
                 norm=False,
                 allow_instantaneous=False):
        self.B = B
        self.dt = dt
        self.dt_max = dt_max
        self.orth = orth
        self.norm = norm
        self.allow_instantaneous = allow_instantaneous

        self.basis = self.interpolate_basis(self.create_basis(), self.dt, self.dt_max, self.norm)
        self.L = self.basis.shape[0]

    @abc.abstractmethod
    def create_basis(self):
        raise NotImplementedError()

    def convolve_with_basis(self, S):
        """
        Convolve each column of the event count matrix with this basis
        :param S:     TxK matrix of inputs.
                      T is the number of time bins
                      K is the number of input dimensions.
        :return: TxKxB tensor of inputs convolved with bases
        """
        (T,K) = S.shape
        (R,B) = self.basis.shape

        # Initialize array for filtered stimulus
        F = np.empty((T,K,B))

        # Compute convolutions fo each basis vector, one at a time
        for b in np.arange(B):
            F[:,:,b] = sig.fftconvolve(S,
                                       np.reshape(self.basis[:,b],(R,1)),
                                       'full')[:T,:]

        # Check for positivity
        if np.amin(self.basis) >= 0 and np.amin(S) >= 0:
            np.clip(F, 0, np.inf, out=F)
            assert np.amin(F) >= 0, "convolution should be >= 0"

        return F

    def interpolate_basis(self, basis, dt, dt_max,
                          norm=True):
        # Interpolate basis at the resolution of the data
        L,B = basis.shape
        t_int = np.arange(0.0, dt_max, step=dt)
        t_bas = np.linspace(0.0, dt_max, L)

        ibasis = np.zeros((len(t_int), B))
        for b in np.arange(B):
            ibasis[:,b] = np.interp(t_int, t_bas, basis[:,b])

        # Normalize so that the interpolated basis has volume 1
        if norm:
            # ibasis /= np.trapz(ibasis,t_int,axis=0)
            ibasis /= (dt * np.sum(ibasis, axis=0))

        if not self.allow_instantaneous:
            # Typically, the impulse responses are applied to times
            # (t+1:t+R). That means we need to prepend a row of zeros to make
            # sure the basis remains causal
            ibasis = np.vstack((np.zeros((1,B)), ibasis))

        return ibasis

    def create_basis(self):
        raise NotImplementedError("Override this in base class")

class CosineBasis(Basis):
    """
    Create a basis of raised cosine tuning curves
    """
    def __init__(self,
                 B, dt, dt_max,
                 orth=False,
                 norm=True,
                 allow_instantaneous=False,
                 n_eye=0,
                 a=1.0/120,
                 b=0.5,
                 L=100):

        self.n_eye = n_eye
        self.a = a
        self.b = b
        self.L = L

        super(CosineBasis, self).__init__(B, dt, dt_max, orth, norm, allow_instantaneous)

    def create_basis(self):
        n_pts = self.L              # Number of points at which to evaluate the basis
        n_cos = self.B - self.n_eye # Number of cosine basis functions'
        n_eye = self.n_eye          # Number of identity basis functions
        assert n_cos >= 0 and n_eye >= 0

        n_bas = n_eye + n_cos
        basis = np.zeros((n_pts,n_bas))

        # The first n_eye basis elements are identity vectors in the first time bins
        basis[:n_eye,:n_eye] = np.eye(n_eye)

        # The remaining basis elements are raised cosine functions with peaks
        # logarithmically warped between [n_eye*dt:dt_max].

        a = self.a                          # Scaling in log time
        b = self.b                          # Offset in log time
        nlin = lambda t: np.log(a*t+b)      # Nonlinearity
        u_ir = nlin(np.arange(n_pts))       # Time in log time
        ctrs = u_ir[np.floor(np.linspace(n_eye,(n_pts/2.0),n_cos)).astype(np.int)]
        if len(ctrs) == 1:
            w = ctrs/2
        else:
            w = (ctrs[-1]-ctrs[0])/(n_cos-1)    # Width of the cosine tuning curves

        # Basis function is a raised cosine centered at c with width w
        basis_fn = lambda u,c,w: (np.cos(np.maximum(-np.pi,np.minimum(np.pi,(u-c)*np.pi/w/2.0)))+1)/2.0
        for i in np.arange(n_cos):
            basis[:,n_eye+i] = basis_fn(u_ir,ctrs[i],w)

        # Orthonormalize basis (this may decrease the number of effective basis vectors)
        if self.orth:
            basis = scipy.linalg.orth(basis)
        if self.norm:
            # We can only normalize nonnegative bases
            if np.any(basis<0):
                raise Exception("We can only normalize nonnegative impulse responses!")

            # Normalize such that \int_0^1 b(t) dt = 1
            basis = basis / np.tile(np.sum(basis,axis=0), [n_pts,1]) / (1.0/n_pts)

        return basis

class IdentityBasis(Basis):

    def __init__(self,
                 dt, dt_max,
                 norm=True,
                 allow_instantaneous=False):

        self.dt = dt
        self.dt_max = dt_max
        self.allow_instantaneous = allow_instantaneous
        self.norm = norm
        self.B = int(dt_max // dt)

        if allow_instantaneous:
            self.B += 1

        # The length is always the current time bin + dt_max / dt
        self.L = dt_max // dt + 1

        # Save allow instantaneous
        self.basis = self.create_basis()

    def create_basis(self):
        ibasis = np.eye(self.B)

        # Typically we don't allow the instantaneous connection
        # since this is no longer a generative model
        if not self.allow_instantaneous:
            ibasis = np.vstack((np.zeros((1,self.B)), ibasis))

        assert ibasis.shape == (self.L, self.B)

        # Normalize so that the interpolated basis has volume 1
        if self.norm:
            ibasis /= (self.dt * np.sum(ibasis, axis=0))

        return ibasis


from classes.LBP import LBP

class Buyers:
    pass

class LimitOrders(Buyers):

    def __init__(self, target: float, budget: float):
        """Trader that buy an amount of token when its price is lower thank or equal to to target.

        Args:
            target: maximum price allowed to buy.
            budget: budget for the buy.
        """
        self.target = target
        self.budget = budget
        self.bought = False

    def buy(self, lbp: LBP):
        """Buy token with amount if price is less than target price and if not a previous buy exists."""
        if not self.bought and self.target >= lbp.spot_price():
            lbp.buy_tokens(usd=self.budget)
            self.bought = True

class FrontRunners(Buyers):

    def __init__(self, budget: float):
        self.budget = budget
        self.bought = False

    def buy(self, lbp):
        if not self.bought:
            lbp.buy_tokens(usd=self.budget)
            self.bought = True

class BuyersFactory:

    @staticmethod
    def build_buyers(buyers_config_list: list) -> list:
        """

        Args:
            buyers_config_list: a list of child of the parent class Buyers

        Returns:

        """
        buyers = []  # this is a queue of buyers' orders
        # If config is an empty list, buyers will be an empty list. In this case the price will be driven by te weights.
        for buyer in buyers_config_list:
            buyer_type = buyer["type"]
            if not buyer_type:
                RuntimeError("No buyer type specified in config")
            elif buyer_type == "FrontRunners":
                buyers.append(FrontRunners(buyer["budget"] * buyer["count"]))
            elif buyer_type == "LimitOrders":
                for _ in range(buyer["count"]):
                    buyers.append(LimitOrders(buyer["target"], buyer["budget"]))

        return buyers