from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
from tensorflow.compat.v1 import app
from experiments.adv import run_experiment, run_temperature_scaling

args = argparse.ArgumentParser(description='test malware detectors on the adversarial dataset')
args.add_argument('--detector', type=str, default='drebin', choices=[
                    'drebin',  # deepdrebin
                    'opcodeseq',  # deepDroid
                    'multimodality',  # multimodalnn
                    'dex2img',  # r2d2, due to the issue of effectiveness, we neglect this method
                    'apiseq',  # droidetec, due to the issue of effectiveness, we neglect this method
                ], help='malware detection method')

args.add_argument('--calibration', type=str, default='vanilla', choices=[
                    'vanilla',
                    'temp_scaling'
                    'mc_dropout',
                    'deep_ensemble',
                    'weighted_ensemble',
                    'bayesian'  # variational bayesian inference
                ], help='calibration method')
args.add_argument('--proc_numbers', type=int, default=2,
                  help='number of threads for parallelizing features extraction.')

option = args.parse_args()

non_used_methods = ['dex2img', 'apiseq']
assert option.detector not in non_used_methods


def main(_):
    if option.calibration != 'temp_scaling':
        run_experiment(option.detector,
                       option.calibration,
                       option.proc_numbers
                       )
    else:
        run_temperature_scaling(option.detector, 'vanilla', option.proc_numbers)


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