#! /bin/env python
# https://cms-analysis.github.io/HiggsAnalysis-CombinedLimit/part2/bin-wise-stats/
import os, os.path, sys, stat, argparse, getpass, json
import subprocess
import json
import glob
import ROOT
ROOT.gROOT.SetBatch()
ROOT.PyConfig.IgnoreCommandLineOptions = True
import logging
logger = logging.getLogger(__name__)

from math import sqrt
from datetime import datetime

import numpy as np
import Harvester as H
import Constants as Constants
import CombineHarvester.CombineTools.ch as ch

import logging
LOG_LEVEL = logging.DEBUG
stream = logging.StreamHandler()
stream.setLevel(LOG_LEVEL)
logger = logging.getLogger(__name__)
logger.setLevel(LOG_LEVEL)
logger.addHandler(stream)
try:
    import colorlog
    from colorlog import ColoredFormatter
    formatter = ColoredFormatter(
                "%(log_color)s%(levelname)-8s%(reset)s %(log_color)s%(message)s",
                datefmt=None,
                reset=True,
                log_colors={
                        'DEBUG':    'cyan',
                        'INFO':     'green',
                        'WARNING':  'blue',
                        'ERROR':    'red',
                        'CRITICAL': 'red',
                        },
                secondary_log_colors={},
                style='%')
    stream.setFormatter(formatter)
except ImportError:
    print(' try https://pypi.org/project/colorlog/  for colorlog')
    pass

def get_Nm_for_runmode(mode):
    Nm_formode = {'ellipse':'rho_steps',
                  'dnn'    :'dnn_scores',
                  'mjj'    :'mjj_bins',
                  'mlljj'  :'mlljj_bins',
                  'mjj_vs_mlljj' : 'mjj_vs_mlljj_map',
                  'mjj_and_mlljj': 'mjj_and_mlljj_combined_bins', }
    return Nm_formode[mode]

def format_parameters(p):
    mH = "%.2f" % p[0]
    mA = "%.2f" % p[1]
    return ("MH_"+str(mH) + "_" + "MA_"+str(mA)).replace(".", "p")

def format_ellipse(p, ellipses):
    mH = "%.2f" % p[0]
    mA = "%.2f" % p[1]
    for ie, e in enumerate(ellipses['MuMu'], 0): #put here the enumerate index!!!
        # it does not matter which flavor we pick as we are checking only the ellipse index
        mA_file = "%.2f" % e[-2]
        mH_file = "%.2f" % e[-1]
        if mA == mA_file and mH == mH_file: # check sim masses
            return '{:d}'.format(ie)        # return ellipse index

def parameter_type(s):
    try:
        if s == 'all':
            return 'all'
        x, y = map(float, s.replace("m", "-").split(','))
        if x.is_integer():
            x = int(x)
        if y.is_integer():
            y = int(y)
        return x, y
    except:
        raise argparse.ArgumentTypeError("Parameter must be x,y")

def get_hist_regex(r):
    return '^%s(__.*(up|down))?$' % r

def get_signal_parameters(f):
    # HToZATo2L2B_MH-200_MA-50.root
    split_filename = f.replace('.root','').replace('HToZATo2L2B_','')
    split_filename = split_filename.split('_')
    MH = split_filename[0].split('-')[1]
    MA = split_filename[1].split('-')[1]
    return int(MH), int(MA)

signal_grid = [
        #part0 : 21 signal samples
        ( 200, 50), ( 200, 100), ( 300, 100),
        ( 250, 50), #( 250, 100),
        #( 300, 50),
        ( 300, 100), #( 300, 200),
        #( 500, 50), ( 500, 200), ( 500, 400),
        #( 650, 50),
        #( 800, 50), ( 800, 200), ( 800, 400), ( 800, 700),
        #(1000, 50), (1000, 200), (1000, 500),
        #(2000, 1000),
        #(3000, 2000)
        ]
extra_signals = [
        ]

def prepare_DataCards(grid_data= None, dataset=None, expectSignal=None, era= None, parameters= None, mode= None, input= None, ellipses_mumu_file= None, output= None, method= None, node= None, blind= False, signal_strength= False, stat_only= False, verbose= False, split_by_categories= False, scale= False):

    luminosity = Constants.getLuminosity(era)

    signal_grid = list(set(grid_data))
    parameters = []
    for f in glob.glob(os.path.join(options.input, '*.root')):
        split_filename = f.split('/')[-1]
        if not split_filename.startswith('HToZATo2L2B_'):
            continue
        mH, mA = get_signal_parameters(split_filename)
        if (mH, mA) in signal_grid:
            parameters.append( (mH, mA) )

    if signal_strength:
        parameters = [( 500, 300)]
    if len(parameters) == 1 and parameters[0] == 'all':
        parameters = parameters[:]
    for p in parameters:
        if not p in parameters:
            print("Invalid parameter '%s'. Valid values are %s" % (str(p), str(parameters)))
            return

    print("\tEra and the corresponding luminosity      : %s, %s" %(era, Constants.getLuminosity(era)))
    print("\tInput path                                : %s" %input )
    print("\tGenerating set of cards for parameter(s)  : %s" % (', '.join([str(x) for x in parameters])))
    print("\tChosen analysis mode                      : %s" % mode)

    ellipses = {}
    ellipses['MuMu'] = []
    ellipses['ElEl'] = []
    ellipses['MuEl'] = []
    with open(ellipses_mumu_file.replace('ElEl', 'MuMu')) as inf:
        content = json.load(inf)
        ellipses['MuMu'] = content
        ellipses['MuEl'] = content
    with open(ellipses_mumu_file.replace('MuMu', 'ElEl')) as inf:
        content = json.load(inf)
        ellipses['ElEl'] = content

    prepareShapes(input, dataset, expectSignal, era, method, parameters, ['ggH', 'bbH'], ['boosted', 'resolved'], ['MuMu', 'ElEl'], ellipses, mode, output, luminosity, split_by_categories, scale, blind, signal_strength, stat_only, verbose)

    # Create helper script to run limits
    output = os.path.join(output, mode)
    print( '\tThe generated script to run limits can be found in : %s/' %output)
    script = """#! /bin/bash
scripts=`find {output} -name "*_run_{method}.sh"`
for script in $scripts; do
    dir=$(dirname $script)
    script=$(basename $script)
    echo "\tComputing with ${{script}}"
    pushd $dir &> /dev/null
    . $script
    popd &> /dev/null
done
""".format(output=output, method=method)

    script_name = "run_combined_%s_%s.sh" % (mode, method)
    with open(script_name, 'w') as f:
        f.write(script)

    st = os.stat(script_name)
    os.chmod(script_name, st.st_mode | stat.S_IEXEC)

    if method=="hybridnew":
        logger.info("All done. You can run everything by executing %r" % ('./' + script_name[:-3]+"_onSlurm.sh"))
    else:
        logger.info("All done. You can run everything by executing %r" % ('./' + script_name))

def prepareShapes(input=None, dataset=None, expectSignal=None, era=None, method=None, parameters=None, productions=None, regions=None, flavors=None, ellipses=None, mode=None, output=None, luminosity=None, split_by_categories=False, scale=False, blind=False, signal_strength=False, stat_only=False, verbose=False):
    if mode == "mjj_and_mlljj":
        categories = [
                (1, 'mlljj'),
                (2, 'mjj')
                ]
    elif mode == "mjj_vs_mlljj":
        categories = [
                (1, 'mjj_vs_mlljj')
                ]
    elif mode == "mjj":
        categories = [
                (1, 'mjj')
                ]
    elif mode == "mlljj":
        categories = [
                (1, 'mlljj')
                ]
    elif mode == "ellipse":
        categories = [
                # % (flavour, ellipse_index)
                (1, 'ellipse_{}_{}')
                ]
    elif mode == "postfit":
        split_by_categories = True
        categories = [
                # % ( MH, MA)
                (1, 'dnn_ggH_resolved_MH_{}_MA_{}')
                # % (flavour, ellipse_index)
                #(2, 'ellipse_{}_{}'),
                #(3, 'mlljj'),
                #(4, 'mjj'),
                #(5, 'mjj_vs_mlljj')
                ]
    elif mode == "dnn":
        categories = [
                # % ( MH, MA)
                (1, 'dnn_ggH_resolved_MH_{}_MA_{}')
                ]

    histfactory_to_combine_categories = {}
    histfactory_to_combine_processes  = {
            # main Background
            'ttbar'    : ['^TT*', '^ttbar_*'],
            'SingleTop': ['^ST_*'],
            'DY'       : ['^DYJetsToLL_0J*', '^DYJetsToLL_1J*', '^DYJetsToLL_2J*', '^DYToLL_*'],
            # Others Backgrounds
            'WPlusJets': ['^WJetsToLNu*'],
            'ttV'      : ['^TT(WJets|Z)To*'],
            'VV'       : ['^(ZZ|WW|WZ)To*'],
            'VVV'      : ['^(ZZZ|WWW|WZZ|WWZ)*'],
            #'Wgamma'  : ['^WGToLNuG_TuneCUETP8M1'], TODO add this sample
            'SMHiggs'  : ['^ggZH_HToBB_ZToNuNu*', '^HZJ_HToWW*', '^ZH_HToBB_ZToLL*', '^ggZH_HToBB_ZToLL*', '^ttHJet*']
            }
    if options.mode == "postfit":
        for p in parameters:
            mH = p[0]
            mA = p[1]

            formatted_p = format_parameters(p)
            formatted_e = format_ellipse(p, ellipses)

            histfactory_to_combine_categories = {
                'dnn_ggH_resolved_MH_{}_MA_{}'.format(mH, mA) : get_hist_regex('DNNOutput_{flavor}channel_resolvedselection_ZAscan_MA_%s_MH_%s'%(mA, mH)),
                #'ellipse_{}_{}'.format(formatted_p, formatted_e):    get_hist_regex('rho_steps_histo_{flavor}_hZA_lljj_deepCSV_btagM_mll_and_met_cut_%s' % format_ellipse(p, ellipses)),
                #'mjj':      get_hist_regex('jj_M_{flavor}_hZA_lljj_deepCSV_btagM_'),
                #'mlljj': get_hist_regex('lljj_M_{flavor}_hZA_lljj_deepCSV_btagM_'),
            }
    # Shape depending on the signal hypothesis
    for p in parameters:
        mH = p[0]
        mA = p[1]

        formatted_p = format_parameters(p)
        formatted_e = format_ellipse(p, ellipses)

        suffix = formatted_p.replace('MH_', 'MH-').replace('MA_','MA-')
        # Signal process
        histfactory_to_combine_processes['HToZATo2L2B_MH-%s_MA-%s'%(mH,mA), p] = ['^HToZATo2L2B_MH-%s_MA-%s*'%(mH, mA)]

        if mode == "postfit":
            histfactory_to_combine_categories[('dnn_ggH_resolved_MH_{}_MA_{}'.format(mH, mA), p)] = get_hist_regex('DNNOutput_{flavor}channel_resolvedselection_ZAscan_MA_%s_MH_%s'%(mA, mH))
        elif mode == "mjj_and_mlljj":
            histfactory_to_combine_categories[('mjj', p)]   = get_hist_regex('jj_M_resolved_{flavor}_hZA_lljj_DeepCSVM_mll_and_met_cut')
            histfactory_to_combine_categories[('mlljj', p)] = get_hist_regex('lljj_M_resolved_{flavor}_hZA_lljj_DeepCSVM_mll_and_met_cut')
        elif mode == "mjj_vs_mlljj":
            histfactory_to_combine_categories[('mjj_vs_mlljj', p)] = get_hist_regex('Mjj_vs_Mlljj_resolved_{flavor}_hZA_lljj_DeepCSVM_mll_and_met_cut')
        elif mode == "mlljj":
            histfactory_to_combine_categories[('mlljj', p)] = get_hist_regex('{flavor}_resolved_METCut__mllbb_DeepCSVM')
        elif mode == "mjj":
            histfactory_to_combine_categories[('mjj', p)]   = get_hist_regex('jj_M_resolved_{flavor}_hZA_lljj_DeepCSVM_mll_and_met_cut')
        elif mode == "ellipse":
            histfactory_to_combine_categories[('ellipse_{}_{}'.format(formatted_p, formatted_e), p)] = get_hist_regex('rho_steps_resolved_histo_{flavor}_hZA_lljj_DeepCSV_btagM__METCut__%s'%formatted_p)
        elif mode == "dnn": # FIXME ugly histos name
            #histfactory_to_combine_categories[('dnn_ggH_resolved_MH_{}_MA_{}'.format(mH, mA), p)] = get_hist_regex('DNNOutput_{flavor}channel_resolvedselection_ZAscan_MA_%s_MH_%s'%(mA, mH))
            histfactory_to_combine_categories[('dnn_ggH_resolved_MH_{}_MA_{}'.format(mH, mA), p)] = get_hist_regex('DNNOutput_ZAnode_{flavor}_resolved_DeepCSVM_METCut_gg_fusion_MH_%s_MA_%s'%(mH, mA))

    print( '\thistfactory_to_combine_categories    : %s '%histfactory_to_combine_categories )

    if not blind:
        histfactory_to_combine_processes['data_obs'] = ['^DoubleMuon*', '^DoubleEG*', '^MuonEG*', '^SingleMuon*', '^EGamma*']

    H.splitJECBySources = False
    if signal_strength:
        H.scaleZAToSMCrossSection = True
    H.splitTTbarUncertBinByBin = False

    output_filename = os.path.join(output, 'shapes_HToZATo2L2B.root')
    file, systematics = H.prepareFile(processes_map=histfactory_to_combine_processes, categories_map=histfactory_to_combine_categories, input=input, output_filename=output_filename, signal_process='HToZATo2L2B', method=method, luminosity=luminosity, flavors=flavors, era=era, blind=blind)

    #print ( "\tsystematics : %s       :" %systematics )
    for i, p in enumerate(parameters):
        mH = p[0]
        mA = p[1]

        cb = ch.CombineHarvester()
        if verbose:
            cb.SetVerbosity(3)
        cb.SetFlag("zero-negative-bins-on-import", True)

        # Dummy mass value used for all combine input when a mass is needed
        mass = "125"
        formatted_p = format_parameters(p)
        formatted_e = format_ellipse(p, ellipses)

        analysis_name = 'HToZATo2L2B'
        categories_with_parameters = categories[:]
        for i, k in enumerate(categories_with_parameters):
            if mode in ['dnn', 'postfit']:
                categories_with_parameters[i] = (k[0], k[1].format(mH, mA))
            elif mode == 'ellipse':
                categories_with_parameters[i] = (k[0], k[1].format('MH-%s_MA-%s'%(mH, mA), formatted_e))
            else:
                print( 'FIXME' )
        print( '\tcategories_with_parameters      :  %s '%categories_with_parameters)
        #cb.AddObservations( mass, analysis, era, channel, bin)
        cb.AddObservations(['*'], [analysis_name], ['13TeV_%s'%era], flavors, categories_with_parameters)
        bkg_processes = [
                'ttbar',
                'SingleTop',
                'DY',
                #'WPlusJets',
                #'ttV',
                #'VV',
                #'VVV',
                #'Wgamma',
                #'SMHiggs'
                ]

        # FIXME I am not sure if this intended or a mistake
        for flavor in flavors:
            processes = []
            cb.AddProcesses(['*'], [analysis_name], ['13TeV_%s'%era], [flavor], bkg_processes, categories_with_parameters, False)
            processes += bkg_processes
        sig_process = 'HToZATo2L2B'
        cb.AddProcesses([mass], [analysis_name], ['13TeV_%s'%era], flavors, [sig_process], categories_with_parameters, True)
        processes += [sig_process]
        print ( "\tprocesses       : %s" %processes)

        if not stat_only:
            processes_without_weighted_data = cb.cp()
            processes_without_weighted_data.FilterProcs(lambda p: 'data' in p.process())
            processes_without_weighted_data.AddSyst(cb, 'lumi_$ERA', 'lnN', ch.SystMap('era')(['13TeV_%s'%era], Constants.getLuminosityUncertainty(era)))

            if not H.splitTTbarUncertBinByBin:
                cb.cp().AddSyst(cb, 'ttbar_xsec', 'lnN', ch.SystMap('process')
                        (['ttbar'], 1.001525372691124) )
            cb.cp().AddSyst(cb, 'SingleTop_xsec', 'lnN', ch.SystMap('process')
                    (['SingleTop'], 1.0029650414264797) )
            cb.cp().AddSyst(cb, 'DY_xsec', 'lnN', ch.SystMap('process')
                        (['DY'], 1.007841991384859) )
            if signal_strength:
                # https://twiki.cern.ch/twiki/bin/view/LHCPhysics/LHCHXSWGHH#Current_recommendations_for_di_H
                cb.cp().AddSyst(cb, '$PROCESS_xsec', 'lnN', ch.SystMap('process')
                    ([sig_process], 1.0729) )

            for _, category_with_parameters in categories_with_parameters:
                for flavor in flavors:
                    for process in processes:
                        process = str(process)
                        if sig_process in process:
                            process = sig_process
                        if not process in cb.cp().channel([flavor]).process_set():
                            print("[{}, {}] Process '{}' not found, skipping systematics".format(category_with_parameters, flavor, process))
                        for s in systematics[flavor][category_with_parameters][process]:
                            s = str(s)
                            if H.ignoreSystematic(flavor, process, s):
                                print("[{}, {}, {}] Ignoring systematic '{}'".format(category_with_parameters, flavor, process, s))
                                continue
                            #FIXME s = H.renameSystematic(flavor, process, s)
                            cb.cp().channel([flavor]).process([process]).AddSyst(cb, s, 'shape', ch.SystMap()(1.00))

        # Import shapes from ROOT file
        for flavor in flavors:
            cb.cp().channel([flavor]).backgrounds().ExtractShapes(file, '$BIN/$PROCESS_%s' % flavor, '$BIN/$PROCESS_%s__$SYSTEMATIC' % flavor)
            cb.cp().channel([flavor]).signals().ExtractShapes(file, '$BIN/$PROCESS_MH-%s_MA-%s_%s' % (p[0],p[1], flavor), '$BIN/$PROCESS_%s_%s__$SYSTEMATIC' % ('MH-%s_MA-%s'%(p[0],p[1]), flavor))

        if scale:
            cb.cp().process(['HToZATo2L2B']).ForEachProc(lambda x : x.set_rate(x.rate()*1000))
            cb.cp().process(['HToZATo2L2B']).PrintProcs()

        # Bin by bin uncertainties
        if not stat_only:
            bkgs = cb.cp().backgrounds()
            bkgs.FilterProcs(lambda p: 'data' in p.process())
            bbb = ch.BinByBinFactory()
            bbb.SetAddThreshold(0.05).SetMergeThreshold(0.5).SetFixNorm(False)
            bbb.MergeBinErrors(bkgs)
            bbb.AddBinByBin(bkgs, cb)

        output_prefix      = 'HToZATo2L2B'
        #output_prefix_run = 'HToZATo2L2B_%s' % (formatted_e)

        output_dir = os.path.join(output, mode, H.get_method_group(method), 'MH-%s_MA-%s'%(p[0],p[1]))
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        def createRunCombineScript(mass, output_dir, output_prefix):
            # Write small script to compute the limit
            datacard = os.path.join(output_dir, output_prefix + '.dat')
            workspace_file = os.path.basename(os.path.join(output_dir, output_prefix + '_combine_workspace.root'))

            if method == 'hybridnew':
                script = """#! /bin/bash
pushd {dir}
# If workspace does not exist, create it once
if [ ! -f {workspace_root} ]; then
    text2workspace.py {datacard} -m {mass} -o {workspace_root}
fi
# Run limit
combine {method} --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}.log
combine {method} --expectedFromGrid=0.5 --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}_exp.log
combine {method} --expectedFromGrid=0.84 --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}_P1sigma.log
combine {method} --expectedFromGrid=0.16 --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}_M1sigma.log
combine {method} --expectedFromGrid=0.975 --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}_P2sigma.log
combine {method} --expectedFromGrid=0.025 --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} -S {systematics} &> {name}_M2sigma.log
popd
""".format(workspace_root=workspace_file, datacard=os.path.basename(datacard), name=output_prefix, mass=mass, systematics=(0 if stat_only else 1), method=H.get_combine_method(method), dir=os.path.dirname(os.path.abspath(datacard)))
            else:
                script = """#! /bin/bash
pushd {dir}
# If workspace does not exist, create it once
if [ ! -f {workspace_root} ]; then
    text2workspace.py {datacard} -m {mass} -o {workspace_root}
fi
# Run combined
#combine {method} --X-rtd MINIMIZER_analytic -m {mass} -n {name} {workspace_root} {dataset} --expectSignal {expectSignal} {blind} &> {name}.log
combine {method} -m {mass} -n {name} {workspace_root} {dataset} --expectSignal {expectSignal} {blind} &> {name}.log
popd
""".format(workspace_root=workspace_file, datacard=os.path.basename(datacard), name=output_prefix, mass=mass, systematics=(0 if stat_only else 1), method=H.get_combine_method(method), dir=os.path.dirname(os.path.abspath(datacard)), dataset=('-t -1' if dataset=='asimov' else '-t 8 -s -1'), blind=('--run blind' if blind else ''), expectSignal=expectSignal)

            script_file = os.path.join(output_dir, output_prefix + ('_run_%s.sh' % method))
            print( method, script_file)
            with open(script_file, 'w') as f:
                f.write(script)

            st = os.stat(script_file)
            os.chmod(script_file, st.st_mode | stat.S_IEXEC)

        # Write card
        def writeCard(c, mass, output_dir, output_prefix, script=True):
            datacard = os.path.join(output_dir, output_prefix + '.dat')
            c.cp().mass([mass, "*"]).WriteDatacard(datacard, os.path.join(output_dir, output_prefix + '_shapes.root'))
            if script:
                createRunCombineScript(mass, output_dir, output_prefix)

        print ("Writing datacards!")
        print (categories_with_parameters )
        mergeable_flavors = ['MuMu', 'ElEl']
        if split_by_categories:
            for flavor in flavors:
                for i, cat in enumerate(categories_with_parameters):
                    cat_output_prefix = output_prefix + '_%s_%s' % (flavor, cat[1])
                    writeCard(cb.cp().bin([cat[1]]).channel([flavor]), mass, output_dir, cat_output_prefix, i + 1 == len(categories_with_parameters))

            for i, cat in enumerate(categories_with_parameters):
                if all(x in flavors for x in mergeable_flavors):
                    print("Merging flavors datacards into a single one for {}".format(cat[1]))
                    # Merge all flavors into a single datacards
                    datacards = ["{flavor}={prefix}_{flavor}_{category}.dat".format(prefix=output_prefix, flavor=x, category=cat[1]) for x in mergeable_flavors]
                    args = ['combineCards.py'] + datacards

                    merged_datacard_name = output_prefix + '_' + '_'.join(mergeable_flavors) + '_' + cat[1]
                    merged_datacard = os.path.join(output_dir, merged_datacard_name + '.dat')
                    with open(merged_datacard, 'w') as f:
                        subprocess.check_call(args, cwd=output_dir, stdout=f)

                    createRunCombineScript(mass, output_dir, merged_datacard_name)

            for flavor in flavors:
                if mode == "postfit":
                    name = get_Nm_for_runmode('dnn')
                    script = """#! /bin/bash
pushd {dir}
# Fit the {name} distribution
./{prefix}_{categories}_run_fit.sh

# Create post-fit shapes for all the categories
for CAT in {categories}; do
    text2workspace.py {prefix}_${{CAT}}.dat -m {mass} -o {prefix}_${{CAT}}_combine_workspace.root
    PostFitShapesFromWorkspace -w {prefix}_${{CAT}}_combine_workspace.root -d {prefix}_${{CAT}}.dat -o postfit_shapes_${{CAT}}.root -f fitDiagnostics{prefix}_${{CAT}}.root:fit_b -m {mass} --postfit --sampling --samples 1000 --print
    $CMSSW_BASE/src/ZAStatAnalysis/utils/convertPostfitShapesForPlotIt.py -i postfit_shapes_${{CAT}}.root -o plotIt_{flavor} --signal-process HToZATo2L2B -n {name}
done
popd
""".format(prefix=output_prefix + '_' + flavor, flavor=flavor, mass=125, parameter='MH-%s_MA-%s'%(mH,mA), categories=' '.join([x[1] for x in categories_with_parameters]), dir=os.path.abspath(output_dir), name=name)
                    script_file = os.path.join(output_dir, output_prefix + '_' + flavor + ('_do_postfit.sh'))
                    with open(script_file, 'w') as f:
                        f.write(script)

                    st = os.stat(script_file)
                    os.chmod(script_file, st.st_mode | stat.S_IEXEC)
        else:
            print ( flavor, flavors, cb.cp().channel([flavor]), mass, output_dir, output_prefix + "_" + flavor )

            for flavor in flavors:
                writeCard(cb.cp().channel([flavor]), mass, output_dir, output_prefix + "_" + flavor)
            if all(x in flavors for x in mergeable_flavors):
                print("Merging flavors datacards into a single one")
                # Merge all flavors into a single datacards
                datacards = ["{flavor}={prefix}_{flavor}.dat".format(prefix=output_prefix, flavor=x) for x in mergeable_flavors]
                args = ['combineCards.py'] + datacards

                merged_datacard = os.path.join(output_dir, output_prefix + '_' + '_'.join(mergeable_flavors) + '.dat')
                with open(merged_datacard, 'w') as f:
                    subprocess.check_call(args, cwd=output_dir, stdout=f)

                createRunCombineScript(mass, output_dir, output_prefix + '_' + '_'.join(mergeable_flavors))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create shape datacards ready for combine')
    parser.add_argument('-i', '--input',        action='store', dest='input', type=str, required=True,
                                                help='HistFactory input path: those are the histograms for signal/data/backgrounds that pass through all the following \n'
                                                     'steps: 1/- final selection ( 2l+bjets pass btagging discr cut + met + corrections + etc... ) \n'
                                                     '       2/- do skim  \n'
                                                     '       3/- DNN trained using these skimmed trees \n'
                                                     '       4/- run bamboo to produce your dnn outputs(prefit plots) with all systematics variations using the model you get from training.\n')
    parser.add_argument('-o', '--output',       action='store', dest='output', required=True, default=None,
                                                help='Output directory')
    parser.add_argument('-e', '--era',          action='store', dest='era', required=True, default=None, choices=['2016', '2017', '2018'],
                                                help='You need to pass your era')
    parser.add_argument('-p', '--parameters',   nargs='+', metavar='MH,MA', dest='parameters', type=parameter_type, default=['all'],
                                                help='Parameters list. Use `all` as an alias for all parameters')
    parser.add_argument('-s', '--stat',         action='store_true', dest='stat_only', required=False, default=False,
                                                help='Do not consider systematic uncertainties')
    parser.add_argument('-v', '--verbose',      action='store_true', required=False, default=False,
                                                help='For debugging purposes , you may consider this argument !')
    parser.add_argument('--expectSignal',       action='store', required=True, type=int, choices=[0, 1],
                                                help=' Is this S+B fit B-Only  ? ')
    parser.add_argument('--mode',               action='store', dest='mode', default='dnn', choices=['mjj_vs_mlljj', 'mjj_and_mlljj', 'postfit', 'mjj', 'mlljj', 'ellipse', 'dnn'],
                                                help='Analysis mode')
    parser.add_argument('--node',               action='store', dest='node', default='ZA', choices=['DY', 'TT', 'ZA'],
                                                help='DNN nodes')
    parser.add_argument('--method',             action='store', dest='method', required=True, default=None, choices=['asymptotic', 'hybridnew', 'fit'],
                                                help='Analysis method')
    parser.add_argument('--blind',              action='store_true', dest='blind', required=False,
                                                help='Use fake data instead of real data')
    parser.add_argument('--signal-strength',    action='store_true', dest="signal_strength", required=False, default=False,
                                                help='Put limit on the signal strength instead of the cross-section')
    parser.add_argument('--ellipses-mumu-file', action='store', dest='ellipses_mumu_file', required=False,
                                default='/home/ucl/cp3/kjaffel/bamboodev/ZA_FullAnalysis/bamboo_/ZATools/scripts_ZA/ellipsesScripts/vers20.06.03Inputs/fullEllipseParamWindow_MuMu.json',
                                                help='file containing the ellipses parameters for MuMu (ElEl is assumed to be in the same directory)')
    parser.add_argument('--scale',              action='store_true', dest='scale', required=False, default=False,
                                                help='scale signal rate')
    parser.add_argument('--dataset',            action='store', dest='dataset', choices=['toys', 'asimov'], required=True, default=None,
                                                help='if asimov:\n'
                                                        '-t -1 will produce an Asimov dataset in which statistical fluctuations are suppressed. \n'
                                                     'if toys: \n'
                                                        '-t N with N > 0. Combine will generate N toy datasets from the model and re-run the method once per toy. \n'
                                                        'The seed for the toy generation can be modified with the option -s (use -s -1 for a random seed). \n'
                                                        'The output file will contain one entry in the tree for each of these toys.\n')
    options = parser.parse_args()
    options.mode = options.mode.lower()

    prepare_DataCards(grid_data= signal_grid + extra_signals, dataset= options.dataset, expectSignal=options.expectSignal, era=options.era, parameters=options.parameters, mode=options.mode, input=options.input, ellipses_mumu_file=options.ellipses_mumu_file, output=options.output, method=options.method, node=options.node, blind=options.blind, signal_strength=options.signal_strength, stat_only=options.stat_only, verbose=options.verbose, split_by_categories=True, scale=options.scale)


import io
import unittest.mock as mock

import doxhooks.dataio as dataio
from doxhooks.errors import DoxhooksDataError, DoxhooksDataFileError
from pytest import fixture, mark

from doxhooks_pytest import withraises

class ExampleNonRestorableRepr:
    def __repr__(self):
        return "0"

class BaseTestDataIO:
    dummy_path = None

class TestLoading(BaseTestDataIO):
    @withraises
    def when_loading_literal_data(self, data_string):
        fake_input_file = io.StringIO(data_string)
        with mock.patch(
                "doxhooks.fileio.open_input", autospec=True,
                return_value=fake_input_file):
            return dataio.load_literals(self.dummy_path)

    @mark.parametrize(
        "data_string, data", [
            ("{}", {}),
        ])
    def test_literal_data_is_loaded_from_a_text_file(self, data_string, data):
        loaded_data = self.when_loading_literal_data(data_string)

        # then the data is returned.
        assert loaded_data == data

    @mark.parametrize("bad_data_string", ["{", "set()"])
    def test_loading_bad_literal_data_is_an_error(self, bad_data_string):
        self.when_loading_literal_data(
            bad_data_string, raises=DoxhooksDataFileError)

        assert self.error

class TestSaving(BaseTestDataIO):
    @fixture(autouse=True)
    def _setup_output_file(self, fake_output_file):
        self.file = fake_output_file

    @withraises
    def when_saving_literal_data(self, data):
        with mock.patch(
                "doxhooks.fileio.open_output", autospec=True,
                return_value=self.file):
            dataio.save_literals(self.dummy_path, data)

    @mark.parametrize(
        "data, data_string", [
            ({}, "{}"),
        ])
    def test_literal_data_is_saved_in_a_text_file(self, data, data_string):
        self.when_saving_literal_data(data)

        # then the data is saved.
        assert self.file.contents == data_string

    @mark.parametrize(
        "non_restorable_data", [set(), ExampleNonRestorableRepr()])
    def test_saving_non_restorable_data_is_an_error(self, non_restorable_data):
        self.when_saving_literal_data(
            non_restorable_data, raises=DoxhooksDataError)

        assert self.error


from flask import Flask, g

from db import RedisClient

# __all__ = ['app']#如何理解

app = Flask(__name__)

def get_conn():
    """
    Opens a new redis connection if there is none yet for the
    current application context.
    """
    if not hasattr(g, 'redis_client'):
        g.redis_client = RedisClient()
    return g.redis_client

@app.route('/')
def index():
    return 'Welcome to Proxy Pool System'

@app.route('/get')
def get_proxy():
    """
    Get a proxy
    """
    conn = get_conn()
    return conn.pop()

if __name__ == '__main__':
    app.run()


from pathlib import Path
from typing import Any, Callable, Dict, Optional, Type

from ziggurat.parser import Parser
from ziggurat.visitor import Renderer

class Template:
    transforms: Dict[str, Callable[[Any], Any]] = {
        "upper": str.upper,
        "lower": str.lower,
        "capitalize": str.capitalize,
    }

    def __init__(
        self,
        source: str,
        encoding: str = "utf8",
        parser_cls: Type[Parser] = Parser,
        renderer_cls: Type[Renderer] = Renderer,
    ):
        self.source = Path(source)
        self.renderer_cls = renderer_cls
        with open(source, "r", encoding=encoding) as tmpl:
            self.ast = parser_cls(tmpl.read()).parse()

    def render(self, ctx: Dict[str, Any]) -> str:
        renderer = self.renderer_cls(ctx, self.transforms, self.source.parent)
        self.ast.accept(renderer)
        return renderer.result

def register_transform(func: Callable[[Any], Any], name: Optional[str] = None):
    if name is None:
        name = func.__name__
    Template.transforms[name] = func
    return func


#!/usr/bin/env python3

from operator import mul
from functools import reduce

for test in range(int(input())):
    size = reduce(mul, [int(n) for n in input().split()])
    if size%2 == 1:
        size += 2**0.5 - 1
    print('Scenario #{}:\n{:.2f}\n'.format(test+1, size))


# Create your views here.

from django.http import HttpResponse,HttpResponseNotFound
import random

def gimmeword(request, resource):
    if resource == 'first':
        word = random.choice(["The table", "My uncle", "A book", "The road", "Music"])
    elif resource == 'second':
        word = random.choice(["talks", "knows", "thinks", "works", "runs", "eats"])
    elif resource == 'third':
        word = random.choice(["fast", "low", "far away", "hard", "in pijama", "while dances cha-cha-cha"])
    else:
        return HttpResponseNotFound('Page not found')
    return HttpResponse(word)


import time

import keyboard
import mouse
import pytest

from ...ypackage.model.input import (
    ButtonEvent, InputRecorder, KeyboardEvent, MoveEvent, WheelEvent)

class TestInputRecorder:

    @pytest.fixture(autouse=True)
    def init_recorder(self):
        self.recorder = InputRecorder()

    def test_recoder_controls(self):
        return
        assert self.recorder.record()
        assert self.recorder.stop()
        assert not self.recorder.events

    def test_add_event(self):
        firsttime = time.time()
        secondtime = time.time()

        events = [
            MoveEvent(1, 2, firsttime),
            MoveEvent(1, 5, secondtime)
        ]

        for event in events:
            self.recorder.add_event(event)

        assert self.recorder.events == events

    def test_add_record(self):
        curtime = time.time()
        events = [
            MoveEvent(1, 2, curtime),
            KeyboardEvent(keyboard.KEY_UP, "a", time=curtime - 0.25),
            WheelEvent(1, curtime - 0.5),
            ButtonEvent(1, 2, curtime - 0.75)
        ]

        for event in events:
            self.recorder.add_event(event)
        assert self.recorder.events == events[::-1]

        self.recorder.events = []
        assert not self.recorder.events

        self.recorder.events = events
        assert self.recorder.events == events[::-1]

    def test_play_records(self):
        return
        curtime = time.time()

        events = [
            MoveEvent(200, 40, curtime),
            KeyboardEvent(keyboard.KEY_DOWN, "a", time=curtime + 0.1),
            KeyboardEvent(keyboard.KEY_UP, "a", time=curtime + 0.25),
            WheelEvent(-50, curtime + 0.5),
            ButtonEvent(mouse.DOWN, mouse.LEFT, curtime + 0.75),
            ButtonEvent(mouse.UP, mouse.LEFT, curtime + 0.85)
        ]
        self.recorder.events = events
        self.recorder.play()

        # Segmentation fault on Circle CI
        # assert (200, 40) == mouse.get_position()

    def test_wrong_event(self):
        with pytest.raises(ValueError):
            self.recorder.events = "YPackage"

        with pytest.raises(ValueError):
            self.recorder.add_event("YEmreAk")


import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Ellipse
import matplotlib.transforms as transforms

def confidence_ellipse(x, y, ax, n_std=3.0, facecolor='none', **kwargs):
    """
    Create a plot of the covariance confidence ellipse of *x* and *y*.

    Parameters
    ----------
    x, y : array-like, shape (n, )
        Input data.

    ax : matplotlib.axes.Axes
        The axes object to draw the ellipse into.

    n_std : float
        The number of standard deviations to determine the ellipse's radiuses.

    **kwargs
        Forwarded to `~matplotlib.patches.Ellipse`

    Returns
    -------
    matplotlib.patches.Ellipse
    """
    if x.size != y.size:
        raise ValueError("x and y must be the same size")

    cov = np.cov(x, y)
    pearson = cov[0, 1]/np.sqrt(cov[0, 0] * cov[1, 1])
    # Using a special case to obtain the eigenvalues of this
    # two-dimensional dataset.
    ell_radius_x = np.sqrt(1 + pearson)
    ell_radius_y = np.sqrt(1 - pearson)
    ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2,
                      facecolor=facecolor, **kwargs)

    # Calculating the standard deviation of x from
    # the squareroot of the variance and multiplying
    # with the given number of standard deviations.
    scale_x = np.sqrt(cov[0, 0]) * n_std
    mean_x = np.mean(x)

    # calculating the standard deviation of y ...
    scale_y = np.sqrt(cov[1, 1]) * n_std
    mean_y = np.mean(y)

    transf = transforms.Affine2D() \
        .rotate_deg(45) \
        .scale(scale_x, scale_y) \
        .translate(mean_x, mean_y)

    ellipse.set_transform(transf + ax.transData)
    return ax.add_patch(ellipse)

# Excel-Datei mit Beisspieldatensatz in ein Dataframe einlesen
df_raw = pd.read_excel('beispieldatensatz.xlsx')

# Je ein Dataframe für m und eines für w erzeugen
df_m = df_raw.loc[df_raw['Geschlecht'] == 'm']
df_w = df_raw.loc[df_raw['Geschlecht'] == 'w']

# Für m und w nur die Daten in ein Dataframe laden, die Merkmalsausprägungen sind
data_m_groesse = df_m['Körpergröße']
data_m_gewicht = df_m['Körpergewicht']
data_w_groesse = df_w['Körpergröße']
data_w_gewicht = df_w['Körpergewicht']
df_summe_groesse = df_raw['Körpergröße']
df_summe_gewicht = df_raw['Körpergewicht']

# Plot erstellen
bp, ax = plt.subplots(figsize=(5, 4))
scatter = ax.scatter(df_summe_groesse, df_summe_gewicht, c='blue', s=20, zorder=5)

# Ellipse für Standardabweichung erzeugen
confidence_ellipse(df_summe_groesse, df_summe_gewicht, ax, \
                   alpha=0.5, facecolor='pink', edgecolor='purple', zorder=2)

# Achsen und Gitternetz konfigurieren
ax.set(xlabel='Körpergröße [in cm]', ylabel='Körpergewicht [in kg]')
ax.yaxis.grid(True, zorder=1)
ax.xaxis.grid(True, zorder=1)

# Vollständigen Plot anzeigen
plt.show()


from typing import Literal, Optional, Union

from . import xl, func_xltypes
from .func_xltypes import UNUSED, Unused
from .xlerrors import NumExcelError, ValueExcelError

dec = "dec"
Base = Literal[bin, dec, oct, hex]

PERMITTED_DIGITS = {
    bin: set("01"),
    oct: set("01234567"),
    hex: set("0123456789ABCDEFabcdef"),
}

BIT_WIDTHS = {bin: 10, oct: 30, hex: 40}

BASE_NUMBERS = {bin: 2, oct: 8, hex: 16}

BOUNDS = {
    frozenset([bin, oct]): 2 ** 9,
    frozenset([bin, dec]): 2 ** 9,
    frozenset([bin, hex]): 2 ** 9,
    frozenset([oct, dec]): 2 ** 29,
    frozenset([oct, hex]): 2 ** 29,
    frozenset([dec, hex]): 2 ** 39,
}

def handle_places(
        places: Union[Unused, func_xltypes.XlAnything]
) -> Optional[int]:
    if places is UNUSED:
        return None

    if isinstance(places, func_xltypes.Boolean):
        raise ValueExcelError('The `places` argument cannot be a boolean.')

    places = int(places)
    if not (1 <= places <= 10):
        raise NumExcelError('The number of places must be between 1 and 10.')

    return places

def handle_number(number: func_xltypes.XlAnything, origin) -> Union[int, str]:
    if isinstance(number, func_xltypes.Boolean):
        raise ValueExcelError('The number cannot be a boolean.')

    if origin == dec:
        return int(number)

    if isinstance(number, func_xltypes.Blank):
        as_str = "0"

    elif isinstance(number, func_xltypes.Number):
        if number.is_decimal and not number.value.is_integer():
            raise NumExcelError('Number is not an integer.')

        as_str = str(int(number))

    elif isinstance(number, func_xltypes.Text):
        as_str = str(number) if number else "0"

    if len(as_str) > 10:
        raise NumExcelError()

    if set(as_str) - PERMITTED_DIGITS[origin]:
        raise NumExcelError('There are invalid characters in the number.')

    return as_str

def pad_zeroes(string, was_negative, places):
    if places is None:
        return string

    desired_length = len(string) if was_negative else places
    if desired_length < len(string):
        raise NumExcelError(
            'The resulting string is longer than the desired length.')

    return string.zfill(desired_length)

def conversion(number, origin, destination, places):
    if origin == dec:
        # if origin base is dec we do not need to convert from it
        value = number
    else:
        # otherwise, convert to int in the appropriate base.
        as_int = int(number, BASE_NUMBERS[origin])

        # magic:
        mask = 1 << BIT_WIDTHS[origin] - 1
        value = (as_int & ~mask) - (as_int & mask)

    # This is another error-check, but it had to be delayed to here rather
    # than in handle_number because we had to convert it from the origin base
    # first.
    bound = BOUNDS[frozenset([origin, destination])]
    if not (-bound <= value < bound):
        raise NumExcelError('The input number is out of bounds')

    # if the destination base is dec, we are done
    if destination == dec:
        return value

    # magic (handles the 2s-complement-like wrapping behavior):
    was_negative = value < 0
    if was_negative:
        value += 1 << BIT_WIDTHS[destination]

    # Otherwise convert to appropriate string representation (return value is
    # (str, bool)).
    result = destination(value)[2:].upper()
    return pad_zeroes(result, was_negative, places)

def convert_bases(number, origin, destination, places=None):
    # If places is None, means it was not passed to this function, i.e. the
    # corresponding Excel function does not take it as an argument (the
    # ___2DEC functions). If it was passed in, we do error checking via the
    # handle_places function.
    if places is not None:
        places = handle_places(places)

    # Next, we do error checking on the number argument. This function does
    # not do any base-conversion logic, it simply raises exceptions if it
    # encounters invalid inputs.  If the origin base is dec, the result is an
    # integer, otherwise it is a string.
    number = handle_number(number, origin)

    # Now we actually convert from one base to another. The result of this can
    # also be either and integer or a string depending on the destination
    # base. If it is a string, it may also be zero-padded depending on the
    # places argument.
    return conversion(number, origin, destination, places)

@xl.register()
@xl.validate_args
def DEC2BIN(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, dec, bin, places)

@xl.register()
@xl.validate_args
def DEC2OCT(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, dec, oct, places)

@xl.register()
@xl.validate_args
def DEC2HEX(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, dec, hex, places)

# the ___2DEC functions give a number not a string, and they do not take a
# `places` parameter
@xl.register()
@xl.validate_args
def BIN2DEC(
        number: func_xltypes.XlAnything
) -> func_xltypes.XlNumber:
    return convert_bases(number, bin, dec)

@xl.register()
@xl.validate_args
def BIN2OCT(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, bin, oct, places)

@xl.register()
@xl.validate_args
def BIN2HEX(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, bin, hex, places)

@xl.register()
@xl.validate_args
def OCT2DEC(
        number: func_xltypes.XlAnything
) -> func_xltypes.XlNumber:
    return convert_bases(number, oct, dec)

@xl.register()
@xl.validate_args
def OCT2BIN(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, oct, bin, places)

@xl.register()
@xl.validate_args
def OCT2HEX(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, oct, hex, places)

@xl.register()
@xl.validate_args
def HEX2DEC(
        number: func_xltypes.XlAnything
) -> func_xltypes.XlNumber:
    return convert_bases(number, hex, dec)

@xl.register()
@xl.validate_args
def HEX2BIN(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, hex, bin, places)

@xl.register()
@xl.validate_args
def HEX2OCT(
        number: func_xltypes.XlAnything,
        places: func_xltypes.XlAnything = UNUSED
) -> func_xltypes.XlText:
    return convert_bases(number, hex, oct, places)


import sqlite3
import pandas as pd

conn = sqlite3.connect('markov_models.db')

sql = 'select features, percent_return from models'
df = pd.read_sql(sql, conn)

all_features_used = []
for key, value in df.iterrows():
    features = eval(df.loc[key, 'features'])
    all_features_used.extend(features)
    """
    features = sorted(features)
    for i in range(len(features)):

        df.loc[key,'feature_'+str(i)] = features[i]
    """
features_used = list(set(all_features_used))
for feature in features_used:
    df[feature] = 0
    for key, values in df.iterrows():
        if feature in df.loc[key, 'features']:
            df.loc[key, feature] = 1

print(df)

result_features = []
for feature in features_used:

    group = list(df.groupby(by=feature))[1]

    group = group[1]['percent_return']

    result_features.append([feature, len(group), group.mean()])

result_df = pd.DataFrame(result_features, columns = ['features', 'num_in_group', 'means'])
result_df = result_df.sort_values(by=['means'])
print(result_df)
overall_average = result_df['means'].mean()
best_features = sorted( result_df[result_df['means']>overall_average]['features'].values )
print(best_features)
#result_df.to_csv('best_features.csv')


from tkinter import Tk, Frame, StringVar, PhotoImage, Canvas
from tkinter.messagebox import askyesno, showwarning
from tkinter.ttk import Label, Button, Radiobutton, Scrollbar, Frame
# from info_exam import InformationExam
# from center_exam import CenterExam
# from footer_exam import FooterExam
# from state_exam import StateExam
# from question import Answer, convertObject2Question
from PIL import Image, ImageTk
import os
import json
import logging
import threading
import time
import string

logging.basicConfig(level=logging.INFO)
def allIs(value, can_loop):
    '''If all question is mark, it will return True'''
    for i in can_loop:
        if i != value:
            return False
    return True

class Question:

    # mac dinh json data da duoc convert sang object trong python
    def __init__(self,master, index: int, content: str, images=[], answers=[]):
        self.master = master
        self.index = index
        self.content = content
        self.images = images
        self.answers = answers # [Answer objects]
        self.choice = StringVar(master=self.master)
        self.state = 'no'  # mark, sure, no

    def __repr__(self):
        return f'''<Question {self.index}>\ncontent: {self.content}\nanswers: {self.answers}\nchoice: {self.choice.get()}\nstate: {self.state}'''

    def changeState2Mark(self):
        '''change state of question to Mark'''
        self.state = 'mark'

    def changeState2No(self):
        '''change state of question to No'''
        self.state = 'no'

    def changeState2Sure(self):
        '''change state of question to Sure'''
        self.state = 'sure'

    def getAnswer(self):
        '''return a string which contain answer'''
        return self.choice.get()

    def getAmountAnswer(self):
        '''return amount answer of this question'''
        return len(self.answers)

    def answerSelected(self):
        '''trả về Answer object được chọn trong Question, nếu Question chưa được chọn thì trả về None'''
        for answer in self.answers:
            if answer.name == self.choice.get():
                return answer
        return None

class Answer:

    def __init__(self, name, content, image='', **kw):
        self.name = name
        self.content = content
        self.image = image

    def __repr__(self):
        return self.content

def convertObject2Question(master,question_data, index):
    logging.info('Convert data from json to Python object')
    index = index
    content = question_data['question']
    images = question_data['images']
    answers = []
    for i in range(len(question_data['answers'])):
        answer = question_data['answers'][i]
        name = list(string.ascii_uppercase)[i]
        content_ans = answer['content']
        image_ans = answer['image']
        ans = Answer(name=name, content=content_ans, image=image_ans)
        answers.append(ans)
    return Question(master=master,index=index, content=content, images=images, answers=answers)

class StatusButton(Button):

    def __init__(self, master, index, callback,**kw):
        super().__init__(master=master)
        self.master = master
        self.index = index
        self.callback =callback

        self.text_variable = StringVar()
        self.text_variable.set(f'{self.index}')
        self['textvariable'] = self.text_variable

        self.photo = PhotoImage(file='icon/question-24.png')
        self['image'] = self.photo
        self.image = self.photo
        self['compound'] = 'right'
        self['command'] = lambda: self.callback(self.index)

    def __str__(self):
        return f'<StatusButton {self.index}>'

    def changeThisButtonToMark(self):
        self.photo = PhotoImage(file='icon/flag.png')
        self['image'] = self.photo
        self.image = self.photo

    def changeThisButtonToSure(self):
        self.photo = PhotoImage(file='icon/tick-24.png')
        self['image'] = self.photo
        self.image = self.photo

class StateExam(Frame):

    def __init__(self, master, listQuestionObjects:list, callback,**kw):
        super().__init__(master=master)
        self.master = master
        self.listQuestionObjects = listQuestionObjects
        self.callback = callback
        self.setupUI()

    def setupUI(self):
        self.canvas = Canvas(self)

        self.scrollbar = Scrollbar(self)
        self.scrollbar['orient'] = 'vertical'
        self.scrollbar['command'] = self.canvas.yview
        self.scrollbar.pack(side='right', fill='y')

        self.canvas['yscrollcommand'] = self.scrollbar.set
        self.canvas.pack(side='left')

        self.frame = Frame(self.canvas)
        self.frame.bind("<Configure>", self.resizeCanvas)
        self.canvas.create_window((0, 0), window=self.frame, anchor='center')
        self.addStatusButton()

    def resizeCanvas(self, event):
        # width button inside and height is center_exam's height
        self.canvas['scrollregion'] = self.canvas.bbox('all')
        self.canvas['width'] = event.width
        self.canvas['height'] = 150

    def addStatusButton(self):
        self.listStateButton = []
        for index in range(len(self.listQuestionObjects)):
            b = StatusButton(self.frame, index=index, callback=self.callback)
            self.listStateButton.append(b)
            b.pack(fill='x')

class FooterExam(Frame):

    def __init__(self, master,listQuestionObjects: list, **kw):
        super().__init__(master=master, **kw)
        self.master = master
        self.setupUI()

    def setupUI(self):
        self.previousButton = Button(self)
        self.previousButton['text'] = 'Previous'
        self.previousButton['state'] = 'disabled'
        self.previousButton.pack(side='left')

        self.nextButton = Button(self)
        self.nextButton['text'] = 'Next'
        self.nextButton.pack(side='left')

        self.doneButton = Button(self)
        self.doneButton['text'] = 'Done'
        self.doneButton.pack(side='right')

class InformationExam(Frame):

    def __init__(self, master, data, callback, **kw):
        super().__init__(master=master)
        self.master = master
        self.data = data
        self.callback =callback

        self.data['questions']
        self.numberQuestion = len(self.data['questions'])
        self.room = self.data['informations']['room']
        self.date = self.data['informations']['date']
        self.name = self.data['informations']['name']
        self.id = self.data['informations']['id']
        self.subject = self.data['informations']['subject']
        self.time = self.data['informations']['time']
        self.time = int(self.time)
        self.setupUI()

    def setupUI(self):
        self.currentQuestion = Label(self)
        self.currentQuestion['text'] = f'Số câu hỏi: {self.numberQuestion}'
        self.currentQuestion.grid(column=0, row=0)

        self.testRoom = Label(self, text=self.room)
        self.testRoom.grid(column=1, row=0)

        self.testDate = Label(self, text=self.date)
        self.testDate.grid(column=2, row=0)

        self.studentName = Label(self, text=self.name)
        self.studentName.grid(column=3, row=0)

        self.studentId = Label(self, text=self.id)
        self.studentId.grid(column=4, row=0)

        self.testSubject = Label(self, text=self.subject)
        self.testSubject.grid(column=5, row=0)

        self.timeVar = StringVar()
        self.testTime = Label(self, textvariable=self.timeVar)
        self.testTime.grid(column=6, row=0)

    def countdown(self):
        '''Đếm ngược thời gian'''
        while self.time >= 0:
            mins, secs = divmod(self.time, 60)
            self.timeVar.set('{:02d}:{:02d}'.format(mins, secs))
            time.sleep(1)
            self.time -= 1
        self.callback()
        return

class ImageObject(Frame):

    def __init__(self, master, name: str, path: str, width: int, **kw):
        super().__init__(master=master, **kw)
        self.name = name
        self.path = path
        self.width = width
        self.setupUI()

    def setupUI(self):
        self.label_image = Label(self)

        self.image_f = Image.open(self.path)
        self.height = int(self.width*self.image_f.height/self.image_f.width)
        self.image = self.image_f.resize(
            (self.width, self.height), Image.ANTIALIAS)

        self.photo = ImageTk.PhotoImage(self.image)
        self.label_image['image'] = self.photo
        self.label_image.image = self.photo
        self.label_image.pack()

        self.label_title = Label(self)
        self.label_title['text'] = self.name
        self.label_title['background'] = '#27ae60'
        self.label_title['foreground'] = 'white'
        self.label_title.pack(fill='x')

class ImageQuestionContainer(Frame):
    '''this class will contain images of question'''

    def __init__(self, master, images: list, **kw):
        super().__init__(master=master)
        self.master = master
        self.images = images
        self.setupUI()

    def setupUI(self):
        self.canvas = Canvas(self)

        self.scrollbar = Scrollbar(self)
        self.scrollbar['orient'] = 'vertical'
        self.scrollbar['command'] = self.canvas.yview
        self.scrollbar.pack(side='right', fill='y')

        self.canvas['yscrollcommand'] = self.scrollbar.set
        self.canvas.pack(side='left')

        self.frame = Frame(self.canvas)
        self.frame.bind("<Configure>", self.resizeCanvas)
        self.canvas.create_window((0, 0), window=self.frame, anchor='center')
        self.addStatusButton()

    def resizeCanvas(self, event):
        # width button inside and height is center_exam's height
        self.canvas['scrollregion'] = self.canvas.bbox('all')
        self.canvas['width'] = event.width
        self.canvas['height'] = 150

    def addStatusButton(self):
        for image in self.images:
            b = ImageObject(self.frame, name=image['name'], path=image['path'], width=500)
            b.pack(fill='x')

class CenterExam(Frame):
    '''this frame contain questions and render them'''
    # mac dinh data da duoc conver tu json sang object

    def __init__(self, master, listQuestionObjects:list, radioCallback, sureCallback,**kw):
        super().__init__(master=master)
        self.master = master
        self.callback = radioCallback
        self.sureCallback = sureCallback

        # contain Question objects
        self.listQuestionObjects = listQuestionObjects
        self.setupUI()

    def setupUI(self):
        self.render(self.listQuestionObjects[0])

    def render(self, question: Question):
        '''render the question object'''

        logging.info(f'render question {question.index}')

        # delete the previous question if any
        for widget in self.grid_slaves():
            widget.destroy()

        # render question
        self.label_question = Label(self)
        self.label_question['text'] = str(question.index)+'. '+question.content
        self.label_question.grid(row=0, column=0, sticky='w')

        # get list of answers
        self.list_radiobutton = []
        for answer in question.answers:
            ans = Radiobutton(self)
            ans['text'] = answer.content
            ans['compound'] = 'bottom'
            ans['value'] = answer.name
            ans['variable'] = question.choice
            ans['command'] = self.callback
            if os.path.exists(answer.image):
                i = Image.open(answer.image)
                width = 100
                height = int(width*i.height/i.width)
                i = i.resize((width, height), Image.ANTIALIAS)
                photo = ImageTk.PhotoImage(i)
                ans['image'] = photo
                ans.image = photo
            self.list_radiobutton.append(ans)

        # render answers
        for i in range(len(self.list_radiobutton)):
            # +1 because question's row is 0
            self.list_radiobutton[i].grid(row=i+1, column=0, sticky='w')

        # render frame which contains images of question
        if allIs(True, [os.path.exists(i['path']) for i in question.images]):
            self.frame_images_of_question = ImageQuestionContainer(
                self, question.images)
            location_column_frame_images_of_question = len(self.grid_slaves())
            self.frame_images_of_question.grid(
                row=0, column=1, rowspan=location_column_frame_images_of_question)

        self.button_sure = Button(self)
        self.button_sure['text'] = 'Sure this question'
        icon_tick = PhotoImage(file='icon/tick-24.png')
        self.button_sure['image'] = icon_tick
        self.button_sure.image = icon_tick
        self.button_sure['compound'] = 'left'
        self.button_sure['command'] = self.sureCallback

        # determind location of sure button, this button will be end of the frame
        location_row = len(self.grid_slaves())
        self.button_sure.grid(row=location_row, column=0)

        # render the disabled question
        if question.state == 'sure':
            self.label_question['state'] = 'disabled'
            for radiobutton in self.list_radiobutton:
                radiobutton['state'] = 'disabled'
            self.button_sure['state'] = 'disabled'

class MyInfoExamThread(threading.Thread):

    def __init__(self, name, info_exam: InformationExam):
        threading.Thread.__init__(self)
        self.name = name
        self.info_exam = info_exam

    def run(self):
        self.info_exam.countdown()

class MyStateExamThread(threading.Thread):

    def __init__(self, name, state_exam: StateExam):
        threading.Thread.__init__(self)
        self.name = name
        self.state_exam = state_exam

    def run(self):
        self.state_exam.checkQuestionState()

class App():

    def __init__(self, data, center_exam: CenterExam, info_exam: InformationExam, footer_exam: FooterExam, state_exam: StateExam):
        self.data = data
        self.master = Tk()
        self.center_exam = center_exam
        self.info_exam = info_exam
        self.footer_exam = footer_exam
        self.state_exam = state_exam

        self.listQuestionObjects = [convertObject2Question(
            master=self.master, question_data=question, index=index) for index, question in enumerate(self.data['questions'])]
        self.currentQuestion = 0
        self.setupUI()

    def setupUI(self):
        self.master.title('Xtest')
        self.master.iconbitmap('icon/xtest_icon.ico')
        self.master.grid_rowconfigure(0, minsize=50)

        self.center_exam = CenterExam(master=self.master, listQuestionObjects=self.listQuestionObjects, radioCallback=self.changeSateQuestionToMark, sureCallback=self.changeStateQuestionToSure)
        self.center_exam.grid(row=1, column=0)

        self.info_exam = InformationExam(master=self.master,data=self.data, callback=self.timeOut)
        self.info_exam.grid(row=0, column=0, columnspan=2)
        self.infoThread = MyInfoExamThread(name='Thread Information Exam', info_exam=self.info_exam)
        self.infoThread.daemon = True
        self.infoThread.start()

        self.footer_exam = FooterExam(master=self.master, listQuestionObjects=self.listQuestionObjects)
        self.footer_exam.grid(row=2, column=0, columnspan=2)
        self.footer_exam.previousButton['command'] = self.gotoPreviousQuestion
        self.footer_exam.nextButton['command'] = self.gotoNextQuestion
        self.footer_exam.doneButton['command'] = self.doneExam

        self.state_exam = StateExam(master=self.master, listQuestionObjects=self.listQuestionObjects, callback=self.gotoQuestion)

        self.state_exam.grid(row=1, column=1)

    def changeSateQuestionToMark(self):
        logging.info(f'change state question {self.currentQuestion} to mark')
        self.listQuestionObjects[self.currentQuestion].changeState2Mark()
        logging.info(self.listQuestionObjects[self.currentQuestion].state)

        self.state_exam.listStateButton[self.currentQuestion].changeThisButtonToMark()
        text = f'{self.state_exam.listStateButton[self.currentQuestion].index}. {self.listQuestionObjects[self.currentQuestion].choice.get()}'
        self.state_exam.listStateButton[self.currentQuestion].text_variable.set(text)

    def changeStateQuestionToSure(self):
        logging.info(f'change state question {self.currentQuestion} to sure')
        if self.listQuestionObjects[self.currentQuestion].choice.get() == '':
            logging.warning(f'Have no complete the question {self.currentQuestion}')
            showwarning('Cảnh báo','Bạn chưa trả lời câu hỏi này!')
        else:
            if askyesno('Thông tin',f'Chúng tôi sẽ chốt câu hỏi này với đáp án {self.listQuestionObjects[self.currentQuestion].choice.get()}, bạn chắc chứ?'):
                self.listQuestionObjects[self.currentQuestion].changeState2Sure()
                self.state_exam.listStateButton[self.currentQuestion].changeThisButtonToSure()
                self.center_exam.label_question['state'] = 'disabled'
                for btt in self.center_exam.list_radiobutton:
                    btt['state'] = 'disable'
                self.center_exam.button_sure['state'] = 'disabled'

    def getAmountQuestion(self):
        '''Trả về sổ lượng câu hỏi'''
        return len(self.listQuestionObjects)

    def gotoNextQuestion(self):
        logging.info(f'go to next question {self.currentQuestion+1}')
        self.currentQuestion = self.currentQuestion +1
        self.footer_exam.previousButton['state'] = 'normal'
        self.center_exam.render(self.listQuestionObjects[self.currentQuestion])
        if self.currentQuestion == self.getAmountQuestion() -1:
            self.footer_exam.nextButton['state'] = 'disabled'
            self.footer_exam.previousButton['state'] = 'normal'

    def gotoPreviousQuestion(self):
        logging.info(f'go to previous question {self.currentQuestion-1}')
        self.currentQuestion = self.currentQuestion -1
        self.footer_exam.nextButton['state'] = 'normal'
        self.center_exam.render(self.listQuestionObjects[self.currentQuestion])
        if self.currentQuestion==0:
            self.footer_exam.previousButton['state'] = 'disabled'
            self.footer_exam.nextButton['state'] = 'normal'

    def gotoQuestion(self, index):
        logging.info(f'goto the question {index}')
        if self.currentQuestion == index:
            return
        logging.info(f'go to question {index}')
        self.currentQuestion = index
        if self.currentQuestion == 0:
            self.footer_exam.previousButton['state'] = 'disabled'
            self.footer_exam.nextButton['state'] = 'normal'
        elif self.currentQuestion == self.getAmountQuestion()-1:
            self.footer_exam.previousButton['state'] = 'normal'
            self.footer_exam.nextButton['state'] = 'disabled'
        else:
            self.footer_exam.previousButton['state'] = 'normal'
            self.footer_exam.nextButton['state'] = 'normal'
        self.center_exam.render(self.listQuestionObjects[index])

    def doneExam(self):
        logging.info('Done exam')
        if allIs('sure',[question.state for question in self.listQuestionObjects]):
            if askyesno('Thông báo','Bạn thật sự muốn nộp bài chứ?'):
                logging.info(f'{[question.getAnswer() for question in self.listQuestionObjects]}')
                self.master.quit()
        elif [question.state for question in self.listQuestionObjects].count('mark') + [question.state for question in self.listQuestionObjects].count('sure') == self.getAmountQuestion():
            if askyesno('Thông báo','Bạn có muốn chốt tất cả các câu hỏi còn lại và nộp bài không?'):
                logging.info(self.getAllAnswers())
                self.master.quit()
        else:
            unanswered_question = ', '.join([str(question.index) for question in self.listQuestionObjects if question.state == 'no'])
            showwarning('Cảnh báo',f'Bạn chưa trả lời {self.getAmountUnansweredQuestion()} câu hỏi, đó là {unanswered_question}')

    def timeOut(self):
        showwarning('Hết giờ','Đã hết giờ, đáp án của bạn sẽ được gửi lên server, chờ một chút để có được điểm của mình!')
        logging.info(self.getAllAnswers())
        self.master.quit()

    def getAmountUnansweredQuestion(self):
        return [question.state for question in self.listQuestionObjects].count('no')

    def getAmountMarkedQuestion(self):
        return [question.state for question in self.listQuestionObjects].count('mark')

    def getAmountCompletedQuestion(self):
        return [question.state for question in self.listQuestionObjects].count('sure')

    def getAllAnswers(self)->list:
        '''Trả về một list chứa đáp án'''
        return [question.getAnswer() for question in self.listQuestionObjects]

    def run(self):
        self.master.mainloop()

if __name__ == '__main__':
    with open('data_server_send.json', 'r', encoding='utf-8') as f:
        data = json.load(f)

    a = App(data=data, center_exam=CenterExam, info_exam=InformationExam,footer_exam=FooterExam, state_exam=StateExam)
    a.run()


from . import rsync as rsync_checker
from . import scp as scp_checker

import argparse
import glob
import io
import ipaddress
import os
import pprint
import re
import shutil
import subprocess
import sys
import textwrap
import time
import traceback
import yaml

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader

if sys.version_info.major >= 3:
    unicode = lambda x: x

def pretty(thing):
    """Return pretty-printable version."""
    ppthing = pprint.PrettyPrinter(indent=4)
    return ppthing.pformat(thing)

class Error(Exception):
    """authprogs error class."""

    pass

class SSHEnvironmentError(Error):
    """Problem with the SSH server-side environment.

    These error messages are show directly to users, so be
    cautious in what you say.
    """

    pass

class ConfigError(Error):
    """Problem with the authprogs configuration."""

    pass

class CommandRejected(Error):
    """Client command rejected.

    These error messages are show directly to users, so be
    cautious in what you say.
    """

    pass

class InstallError(Error):
    """Problem with the installing an authorized_keys entry."""

    pass

class AuthProgs(object):  # pylint: disable-msg=R0902
    """AuthProgs class"""

    def __init__(
        self,
        logfile=None,
        configfile=None,
        configdir=None,
        debug=False,
        **kwargs
    ):
        """AuthProgs constructor.

        kwargs include:
            authprogs_binary: path to this binary, when creating
                              authorized_keys entries.
                              If not specified, determines from sys.argv[0]
            name: the name of this key, for matching in rules.
        """

        self.debug = debug
        self.logfile = logfile
        self.client_ip = None
        self.silent = kwargs.get('silent')
        if logfile:
            self.logfh = open(logfile, 'a')
        else:
            self.logfh = False

        # Stores result of shutils.which
        # key is program name, val is realpath.
        self.which_cache = {}

        if kwargs.get('authprogs_binary'):
            self.authprogs_binary = kwargs['authprogs_binary']
        else:
            self.authprogs_binary = os.path.abspath(
                os.path.abspath(sys.argv[0])
            )

        self.original_command_string = os.environ.get(
            'SSH_ORIGINAL_COMMAND', ''
        )

        self.original_command_list = self.original_command_string.split()

        self.keyname = kwargs.get('keyname')
        if not self.keyname:
            self.keyname = ''
        if ' ' in self.keyname or '\t' in self.keyname:
            self.log('FATAL: keyname contains space/tabs\n')
            raise Error('--keyname may contain neither spaces nor tabs.')

        self.yamldocs = None
        self.configfile = configfile
        self.configdir = configdir

        self.subvalidators = {}
        self.subvalidators['rsync'] = rsync_checker.RsyncValidator(self)
        self.subvalidators['scp'] = scp_checker.ScpValidator(self)

    def __del__(self):
        if self.logfh:
            self.logfh.close()

    def raise_and_log_error(self, error, message):
        """Raise error, including message and original traceback.

        error: the error to raise
        message: the user-facing error message
        """
        self.log(
            'raising {}, traceback {}\n'.format(error, traceback.format_exc())
        )
        raise error(message)

    def get_client_ip(self):
        """Return the client IP from the environment."""

        if self.client_ip:
            return self.client_ip

        try:
            client = os.environ.get(
                'SSH_CONNECTION', os.environ.get('SSH_CLIENT')
            )
            self.client_ip = client.split()[0]
            self.logdebug('client_ip: {}\n'.format(self.client_ip))
            return self.client_ip
        except:
            raise SSHEnvironmentError(
                'cannot identify the ssh client IP address'
            )

    def logdebug(self, message):
        """Log debugging information."""
        if self.debug:
            self.log(message)

    def log(self, message):
        """Log information."""
        if self.logfh:
            self.logfh.write(message)  # pylint: disable-msg=E1103

    def _glob(self, path):
        """Return glob results of path, or original if no match

        glob.glob returns a list only if there are actual matching files.
        So we always return the original if glob doesn't succeed
        so we can create files that don't yet exist, etc.

        Mocked out in unit tests, silly otherwise.
        """
        results = glob.glob(path)
        return glob.glob(path) or path

    def _os_path_abspath(self, path):
        """Return abspath of path

        Mocked out in unit tests, silly otherwise.
        """
        return os.path.abspath(path)

    def _os_path_expanduser(self, path):
        """Return pathname with ~ expansions.

        Mocked out in unit tests, silly otherwise.
        """
        return os.path.expanduser(path)

    def _os_path_realpath(self, path):
        """Return canonical pathname.

        Mocked out in unit tests, silly otherwise.
        """
        return os.path.realpath(path)

    def globpaths(self, path, expanduser=False, realpath=False):
        """Return paths for given file name after glob expansion.

        If expanduser is set then expand ~ to current user's home dir
        If realpath is set then convert to real path after resolving symlinks
        Always calls abspath to create a path that starts at / and resolve ..

        Perform glob expansion of name and returns list of canonicalised
        file names
        """
        path = self._os_path_abspath(path)

        if expanduser:
            path = self._os_path_expanduser(path)

        if realpath:
            path = self.realpath(path)

        paths = self._glob(path)

        return paths

    def realpath(self, path):
        """Return the real path for the given file path.

        Returns canonicalised file name (eliminate symlinks, handle '..')
        """
        return self._os_path_realpath(path)

    def check_keyname(self, rule):
        """If a key name is specified, verify it is permitted."""

        keynames = rule.get('keynames')
        if not keynames:
            self.logdebug('no keynames requirement.\n')
            return True
        if not isinstance(keynames, list):
            keynames = [keynames]

        if self.keyname in keynames:
            self.logdebug('keyname "{}" matches rule.\n'.format(self.keyname))
            return True
        else:
            self.logdebug(
                'keyname "{}" does not match rule.\n'.format(self.keyname)
            )
            return False

    def check_client_ip(self, rule):
        """If a client IP is specified, verify it is permitted."""

        if not rule.get('from'):
            self.logdebug('no "from" requirement.\n')
            return True

        allow_from = rule.get('from')
        if not isinstance(allow_from, list):
            allow_from = [allow_from]

        def ipnet(addr):
            addr = unicode(addr)
            if addr.lower() in ('*', 'any'):
                addr = '0.0.0.0/0'
            try:
                return ipaddress.ip_network(addr, strict=False)
            except ValueError:
                return None

        allow_from = [ipnet(x) for x in allow_from]
        allow_from = filter(lambda x: x, allow_from)
        client_ip = ipaddress.ip_address(unicode(self.get_client_ip()))

        for allow in allow_from:
            if client_ip in allow:
                self.logdebug('client_ip {} in {}\n'.format(client_ip, allow))
                return True
        self.logdebug('client_ip {} not in {}'.format(client_ip, allow_from))
        return False

    def get_merged_config(self):
        """Get merged config file.

        Returns an open StringIO containing the
        merged config file.
        """
        if self.yamldocs:
            return

        loadfiles = []
        if self.configfile:
            loadfiles.append(self.configfile)

        if self.configdir:
            # Gets list of all non-dotfile files from configdir.
            loadfiles.extend(
                [
                    f
                    for f in [
                        os.path.join(self.configdir, x)
                        for x in os.listdir(self.configdir)
                    ]
                    if os.path.isfile(f)
                    and not os.path.basename(f).startswith('.')
                ]
            )

        merged_configfile = io.StringIO()
        merged_configfile.write('-\n')
        for thefile in loadfiles:
            self.logdebug('reading in config file {}\n'.format(thefile))
            with open(thefile, 'r') as merge:
                merged_configfile.write(merge.read())
            merged_configfile.write('\n-\n')
        merged_configfile.seek(0)
        self.logdebug(
            'merged log file: """\n{}\n"""\n'.format(merged_configfile.read())
        )
        merged_configfile.seek(0)
        return merged_configfile

    def load(self):
        """Load our config, log and raise on error."""
        try:
            merged_configfile = self.get_merged_config()
            self.yamldocs = yaml.load(merged_configfile, Loader=Loader)
            merged_configfile.close()

            # Strip out the top level 'None's we get from concatenation.
            # Functionally not required, but makes dumps cleaner.
            self.yamldocs = [x for x in self.yamldocs if x]
            self.logdebug('parsed_rules:\n{}\n'.format(pretty(self.yamldocs)))

        except (yaml.scanner.ScannerError, yaml.parser.ParserError):
            self.raise_and_log_error(ConfigError, 'error parsing config.')

    def dump_config(self):
        """Pretty print the configuration dict to stdout."""
        yaml_content = self.get_merged_config()
        print(
            'Configuration Source:\n\t{}\n'.format(
                yaml_content.read().replace('\n', '\n\t')
            )
        )
        yaml_content.close()

        try:
            self.load()
            print(
                'Effective Configuration\n\t{}\n'.format(
                    pretty(self.yamldocs).replace('\n', '\n\t')
                )
            )
        except ConfigError:
            sys.stderr.write(
                'config parse error. try running with --logfile=/dev/tty\n'
            )
            raise

    def install_key_data(self, keydata, target):
        """Install the key data into the open file."""

        target.seek(0)
        contents = target.read()
        ssh_opts = 'no-port-forwarding'
        keydata = keydata.strip()
        if keydata in contents:
            raise InstallError(
                'key data already in file - refusing to double-install.\n'
            )
        if '\n' in keydata:
            raise InstallError(
                'install file does not appear to be an ssh pubkey\n'
            )
        command = '{} --run'.format(self.authprogs_binary)
        if self.logfile:
            command += ' --logfile={}'.format(self.logfile)
        if self.keyname:
            command += ' --keyname={}'.format(self.keyname)

        target.write(
            'command="{command}",{ssh_opts} {keydata}\n'.format(
                command=command, keydata=keydata, ssh_opts=ssh_opts
            )
        )

    def install_key(self, keyfile, authorized_keys):
        """Install a key into the authorized_keys file."""

        # Make the directory containing the authorized_keys
        # file, if it doesn't exist. (Typically ~/.ssh).
        # Ignore errors; we'll fail shortly if we can't
        # create the authkeys file.
        try:
            os.makedirs(os.path.dirname(authorized_keys), 0o700)
        except OSError:
            pass

        keydata = open(keyfile).read()
        target_fd = os.open(authorized_keys, os.O_RDWR | os.O_CREAT, 0o600)
        self.install_key_data(keydata, os.fdopen(target_fd, 'w+'))

    def find_match_rsync(self, rule):  # pylint: disable-msg=R0911,R0912
        """Handle rsync commands."""
        return self.subvalidators['rsync'].validate_command(
            self.original_command_list[:], rule
        )

    def _whichbin(self, name):
        """shutil.which wrapper for unit test mocking."""
        return shutil.which(name)

    def valid_binary(self, name, allowed):
        """Determine if name is a valid binary in allowed.

        Does realpath expansion and returns actual pathname.
        """
        # Return from cache if set and allowed
        if name in self.which_cache:
            if self.which_cache[name] in allowed:
                return self.which_cache[name]
            else:
                return
        else:
            self.which_cache[name] = None

        # Find it in our path
        foundbin = self._whichbin(name)

        if not foundbin:
            return

        # Convert to realpath, store, and return
        foundbin = os.path.realpath(foundbin)
        if foundbin not in allowed:
            return

        # Store this successful name in cache
        self.which_cache[name] = foundbin
        return foundbin

    def find_match_scp(self, rule):  # pylint: disable-msg=R0911,R0912
        """Handle scp commands."""
        return self.subvalidators['scp'].validate_command(
            self.original_command_list[:], rule
        )

    def find_match_command(self, rule):
        """Return a matching (possibly munged) command, if found in rule."""

        command_string = rule['command']
        command_list = command_string.split()

        self.logdebug(
            'comparing "{}" to "{}"\n'.format(
                command_list, self.original_command_list
            )
        )
        if rule.get('allow_trailing_args'):
            self.logdebug(
                'allow_trailing_args is true - comparing initial list.\n'
            )
            # Verify the initial arguments are all the same
            if self.original_command_list[: len(command_list)] == command_list:
                self.logdebug('initial list is same\n')
                return {'command': self.original_command_list}
            else:
                self.logdebug('initial list is not same\n')

        elif rule.get('pcre_match'):
            if re.search(command_string, self.original_command_string):
                return {'command': self.original_command_list}

        elif command_list == self.original_command_list:
            return {'command': command_list}

    def find_match(self):
        """Load the config and find a matching rule.

        returns the results of find_match_command, a dict of
        the command and (in the future) other metadata.
        """

        self.load()
        for yamldoc in self.yamldocs:
            self.logdebug('\nchecking rule """{}"""\n'.format(yamldoc))

            if not yamldoc:
                continue

            if not self.check_client_ip(yamldoc):
                # Rejected - Client IP does not match
                continue

            if not self.check_keyname(yamldoc):
                # Rejected - keyname does not match
                continue

            rules = yamldoc.get('allow')
            if not isinstance(rules, list):
                rules = [rules]

            for rule in rules:
                rule_type = rule.get('rule_type', 'command')
                if rule_type == 'command':
                    sub = self.find_match_command
                elif rule_type == 'scp':
                    sub = self.find_match_scp
                elif rule_type == 'rsync':
                    sub = self.find_match_rsync
                else:
                    self.log(
                        'fatal: no such rule_type "{}"\n'.format(rule_type)
                    )
                    self.raise_and_log_error(
                        ConfigError, 'error parsing config.'
                    )

                match = sub(rule)
                if match:
                    return match

        # No matches, time to give up.
        if not self.silent:
            sys.stderr.write(
                'command "{}" rejected.\n'.format(self.original_command_string)
            )
        raise CommandRejected(
            'command "{}" denied.'.format(self.original_command_string)
        )

    def exec_command(self):
        """Glean the command to run and exec.

        On problems, sys.exit.
        This method should *never* return.
        """
        if not self.original_command_string:
            raise SSHEnvironmentError(
                'no SSH command found; interactive shell disallowed.'
            )

        command_info = {
            'from': self.get_client_ip(),
            'keyname': self.keyname,
            'ssh_original_command': self.original_command_string,
            'time': time.time(),
        }

        os.environ['AUTHPROGS_KEYNAME'] = self.keyname

        retcode = 126
        try:
            match = self.find_match()
            command_info['command'] = match.get('command')
            self.logdebug('find_match returned "{}"\n'.format(match))

            command = match['command']
            retcode = subprocess.call(command)
            command_info['code'] = retcode
            self.log('result: {}\n'.format(command_info))
            sys.exit(retcode)
        except (CommandRejected, OSError) as err:
            command_info['exception'] = '{}'.format(err)
            self.log('result: {}\n'.format(command_info))
            sys.exit(retcode)

def main():  # pylint: disable-msg=R0912,R0915
    """Main."""
    parser = argparse.ArgumentParser()
    parser.usage = textwrap.dedent(
        """\
    %(prog)s {--run|--install_key|--dump_config} [options]

    SSH command authenticator.

    Used to restrict which commands can be run via trusted SSH keys.
    """
    )

    group = parser.add_argument_group(
        'Run Mode Options',
        'These options determine in which mode the authprogs program runs.',
    )
    group.add_argument(
        '-r',
        '--run',
        dest='run',
        action='store_true',
        help='Act as ssh command authenticator. Use this '
        'when calling from authorized_keys.',
    )
    group.add_argument(
        '--dump_config',
        dest='dump_config',
        action='store_true',
        help='Dump configuration (python format) to standard out and exit.',
    )
    group.add_argument(
        '--install_key',
        dest='install_key',
        help='Install the named ssh public key file to authorized_keys.',
        metavar='FILE',
    )

    group = parser.add_argument_group('Other Options')
    group.add_argument(
        '--keyname',
        dest='keyname',
        help='Name for this key, used when matching config blocks.',
    )
    group.add_argument(
        '--configfile',
        dest='configfile',
        help='Path to authprogs configuration file. '
        'Defaults to ~/.ssh/authprogs.yaml',
        metavar='FILE',
    )
    group.add_argument(
        '--configdir',
        dest='configdir',
        help='Path to authprogs configuration directory. '
        'Defaults to ~/.ssh/authprogs.d',
        metavar='DIR',
    )
    group.add_argument(
        '--logfile',
        dest='logfile',
        help='Write logging info to this file. Defaults to no logging.',
        metavar='FILE',
    )
    group.add_argument(
        '--debug',
        dest='debug',
        action='store_true',
        help='Write additional debugging information to --logfile',
    )
    group.add_argument(
        '--silent',
        dest='silent',
        action='store_true',
        help='Do not tell user their command was rejected on failure.',
    )
    group.add_argument(
        '--authorized_keys',
        dest='authorized_keys',
        default=os.path.expanduser('~/.ssh/authorized_keys'),
        help='Location of authorized_keys file for '
        '--install_key. Defaults to ~/.ssh/authorized_keys',
        metavar='FILE',
    )
    group.add_argument(
        '--version', action='store_true', help='Show version and exit.'
    )

    args = parser.parse_args()

    if args.version:
        from . import __version__

        print(__version__)
        sys.exit()

    if not args.configfile:
        cfg = os.path.expanduser('~/.ssh/authprogs.yaml')
        if os.path.isfile(cfg):
            args.configfile = cfg
    if not args.configdir:
        cfg = os.path.expanduser('~/.ssh/authprogs.d')
        if os.path.isdir(cfg):
            args.configdir = cfg

    if args.debug and not args.logfile:
        parser.error('--debug requires use of --logfile')

    ap = None
    try:
        ap = AuthProgs(
            logfile=args.logfile,  # pylint: disable-msg=C0103
            configfile=args.configfile,
            configdir=args.configdir,
            debug=args.debug,
            keyname=args.keyname,
            silent=args.silent,
        )

        if args.dump_config:
            ap.dump_config()
            sys.exit(0)

        elif args.install_key:
            try:
                ap.install_key(args.install_key, args.authorized_keys)
                sys.stderr.write('Key installed successfully.\n')
                sys.exit(0)
            except InstallError as err:
                sys.stderr.write('Key install failed: {}'.format(err))
                sys.exit(1)

        elif args.run:
            ap.exec_command()
            sys.exit('authprogs command returned - should never happen.')
        else:
            parser.error('Not sure what to do. Consider --help')

    except SSHEnvironmentError as err:
        ap.log(
            'SSHEnvironmentError "{}"\n{}\n'.format(err, traceback.format_exc())
        )
        sys.exit('authprogs: {}'.format(err))
    except ConfigError as err:
        ap.log('ConfigError "{}"\n{}\n'.format(err, traceback.format_exc()))
        sys.exit('authprogs: {}'.format(err))
    except CommandRejected as err:
        sys.exit('authprogs: {}'.format(err))
    except Exception as err:
        if ap:
            ap.log(
                'Unexpected exception: {}\n{}\n'.format(
                    err, traceback.format_exc()
                )
            )
        else:
            sys.stderr.write(
                'Unexpected exception: {}\n{}\n'.format(
                    err, traceback.format_exc()
                )
            )
        sys.exit('authprogs experienced an unexpected exception.')

if __name__ == '__main__':
    sys.exit('This is a library only.')


from cryptography.fernet import Fernet
import argparse

def banner():
    print('''
+-----------------------------------------------+
| magicave file encryptor                       |
| magicave v.0.1                                |
| Author: Elvin --> 'elvinsl'                   |
| Github: https://github.com/elvinsl/magicave   |
+-----------------------------------------------+
    ''')

class EncDec:
    @classmethod
    def generate(cls, file):
        f = Fernet.generate_key()
        with open(file, 'wb') as key_file:
            key_file.write(f)

    @classmethod
    def encrypt(cls, file, key_file):
        key = open(key_file, 'rb').read()
        f = Fernet(key)
        with open(file, 'rb') as normal_file:
            file_data = normal_file.read()
        enc_data = f.encrypt(file_data)
        with open(file, 'wb') as normal_file:
            normal_file.write(enc_data)

    @classmethod
    def decrypt(cls, file, key_file):
        key = open(key_file, 'rb').read()
        f = Fernet(key)
        with open(file, 'rb') as enc_file:
            enc_data = enc_file.read()
        dec_data = f.decrypt(enc_data)
        with open(file, 'wb') as enc_file:
            enc_file.write(dec_data)

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--file', dest='file',
                        help='File to encrypt or decryp')
    parser.add_argument('-k', '--key', dest='key',
                        help='Key to encrypt or decrypt')
    parser.add_argument('-e', '--encrypt', dest='enc',
                        help='Encrypt the file with given key', action='store_true')
    parser.add_argument('-d', '--decrypt', dest='dec',
                        help='Decrypt the file with given key', action='store_true')
    parser.add_argument('-g', '--generate', dest='generate',
                        help='Generate new key file')
    dest_value = parser.parse_args()
    return dest_value

args = get_args()
banner()

if args.generate:
    EncDec().generate(args.generate)
elif args.dec and args.file and args.key:
    EncDec().decrypt(args.file, args.key)
elif args.enc and args.file and args.key:
    EncDec().encrypt(args.file, args.key)
else:
    print('''
Usage: '--help' for more info
Generate Key: python3 magicave.py --generate key.txt
Encrypt File: python3 magicave.py --encrypt --file <file> --key <key>
Encrypt File: python3 magicave.py --decrypt --file <file> --key <key>
    ''')


import obspy
import os
from obspy.core import UTCDateTime
from obspy.io.sac import sactrace

sac_dir=input("PLEASE, enter name of directory containing your sac files:")

sac_files=[os.path.join(sac_dir,sf) for sf in os.listdir(sac_dir) if os.path.isfile(os.path.join(sac_dir,sf))]
catalog=obspy.core.event.Catalog()
#print(sac_files)
#st=obspy.read(pathname_or_url=sfiles)#,format='SAC') #read sac_file to obspy stream
#leave only sac trace files in the file lise:
for sf in sac_files:
   try:
     st=obspy.read(sf,format='SAC')
   except obspy.io.sac.util.SacIOError:
     print('File ',sf,'is not sac trace file, probably response file. Skipping.')
     sac_files.remove(sf)

#read sac headers
for sf in sac_files:
   sac_head=sactrace.SACTrace.read(sf,headonly=True) #sac header
   #extract event info
   if len(catalog)==0: #catalog is empty
    ev_name=sac_head.kevnm #sac event name as resource_id
    #generate Origin from
    #sac_head.o
    #sac_head.evla
    #sac_head.evlo
    #sac_head.evdp
    Origin=obspy.core.event.origin.Origin(time=UTCDateTime(sac_head.reftime+sac_head.o), longitude=sac_head.evlo, latitude=sac_head.evla, depth=sac_head.evdp)
    Magnitude=obspy.core.event.magnitude.Magnitude(mag=sac_head.mag)
    event=obspy.core.event.Event(resource_id=obspy.core.event.ResourceIdentifier(id=ev_name), origins=[Origin],magnitudes=[Magnitude])
    catalog=obspy.core.event.Catalog(events=[event])
   #check if event exists...
   else:
    if not (UTCDateTime(sac_head.reftime+sac_head.o) in [e.origins[0].time for e in catalog] and
       sac_head.evlo in [e.origins[0].longitude for e in catalog]   and
       sac_head.evla in [e.origins[0].latitude for e in catalog] and
       sac_head.evdp in [e.origins[0].depth for e in catalog] and
       sac_head.mag in [e.magnitudes[0].mag for e in catalog]):
    #if not (orig==Origin and magn==Magnitude):     not working probably due to resource_id....
      orig=obspy.core.event.origin.Origin(time=UTCDateTime(sac_head.reftime+sac_head.o), longitude=sac_head.evlo, latitude=sac_head.evla, depth=sac_head.evdp)
      magn=obspy.core.event.magnitude.Magnitude(mag=sac_head.mag)
      ev2=obspy.core.event.Event(resource_id=obspy.core.event.ResourceIdentifier(id=sac_head.kevnm), origins=[orig],magnitudes=[magn])
      catalog.append(ev2)

print('Found ',len(catalog),' events in sac files in directory: ',sac_dir)

#extract picks and add to first event in catalog #WARNING FOR CATALOGS LONGER THAN 1 THIS IS OMITTED
print('Obtaining picks:')
if len(catalog)==1:
   event=catalog[0]
   picks=[]
   for sf in sac_files:
       sac_head=sactrace.SACTrace.read(sf,headonly=True) #read sac header
       wfid=obspy.core.event.WaveformStreamID(network_code=sac_head.knetwk,station_code=sac_head.kstnm,channel_code=sac_head.kcmpnm) #prepare waveform_id for stream identification
       if not sac_head.a==None: #if not empty -> associate as P pick'
          print('Adding pick found in sac.header.a as P from ',wfid.get_seed_string())
          pick=obspy.core.event.origin.Pick(time= sac_head.reftime+sac_head.a, waveform_id=wfid,phase_hint='P')
          picks.append(pick)
       if not sac_head.t0==None: #if not empty -> associate as S pick'
       	  print('Adding pick found in sac.header.t0 as S from ',wfid.get_seed_string())
          pick=obspy.core.event.origin.Pick(time= sac_head.reftime+sac_head.t0, waveform_id=wfid,phase_hint='S')
          picks.append(pick)
   event.picks=picks #append picks to event
   print('Adding nodal planes to event:')
   print(event.origins[0])
   #input in str,dp,rk on 1 line
   str1, dip1, rake1=input('PLEASE, enter strike,dip,rake of nodal plane 1:').split()
   str2, dip2, rake2=input('PLEASE, enter strike,dip,rake of nodal plane 2:').split()
   nodals1=obspy.core.event.source.NodalPlane(strike=str1,dip=dip1,rake=rake1) #generate NP object for first NP
   nodals2=obspy.core.event.source.NodalPlane(strike=str2,dip=dip2,rake=rake2) #generate NP object for second NP

   event.focal_mechanisms=[obspy.core.event.source.FocalMechanism(nodal_planes=obspy.core.event.source.NodalPlanes(nodal_plane_1=nodals1,nodal_plane_2=nodals2))] #associate NP info to focal mechanism of event
   catalog[0]=event #rewrite in catalog with added info
#print(event)
#Write quake ML file
print('Writing QUAKEML file')
outfile=input('please, enter name of quakeml file:')#'quake.ml'
if outfile=='': #if no file name from user use default:
  outfile='quake.ml'
catalog.write(filename=outfile,format='QUAKEML')