<reponame>L1NNA/malware_rl
import os
import sys
import logging
import argparse

##TODO: malconv.h5 that is beign saved is not specified
def research(agent="acer", dynamic = False, rounds = 100, base_dir="sets", gpu_id=0, train=False):

    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
    os.environ["CUDA_VISIBLE_DEVICES"]=str(gpu_id)
    
    from h4rm0ny.train_agent_chainer import train_agent, create_acer_agent, create_ddqn_agent
    from h4rm0ny.envs.utils import malconv
    from h4rm0ny import init
    init()
    
    if dynamic:
        base_dir = os.path.join(base_dir, f'{agent}-dynamic-{rounds}')
    else:
        base_dir = os.path.join(base_dir, f'{agent}-static-{rounds}')
    os.makedirs(base_dir, exist_ok=True)
    test_set = os.path.join(base_dir, "test")
    train_set =os.path.join(base_dir, "train")
    os.makedirs(test_set, exist_ok=True)
    os.makedirs(train_set, exist_ok=True)
    malconv.model_path_uptrain = os.path.join(base_dir, 'malconv_uptrained.h5')
    
    if dynamic:
        if agent == 'acer':
            gym_env = 'malconv-uptrain-v0'
        else:
            gym_env = 'malconv-uptrain-v1'
        direc = 'uptrain'
    else:
        if agent == 'acer':
            gym_env = 'malconv-train-v0'
        else:
            gym_env = 'malconv-train-v1'
        direc= 'no_uptrain'
    if agent == "random":
        agent_score = train_agent(rounds, use_score=True, name="random", create_agent=create_acer_agent,test_set=test_set, train_set=train_set, gym_env = gym_env, train=train )         
    if agent == "acer":
        agent_score = train_agent(rounds, use_score=True, name=agent+'/'+direc, create_agent=create_acer_agent,test_set=test_set, train_set=train_set, gym_env = gym_env, train = train )         
    if agent == "ddqn":
        agent_score = train_agent(rounds, use_score=True, name=agent+'/'+direc, create_agent=create_ddqn_agent, test_set=test_set, train_set=train_set ,gym_env = gym_env, train=train)    
    if dynamic:  
        print("Malconv Model:", malconv.model_path_uptrain)
    else:
        print("Malconv Model: h4rm0ny/envs/utils/malconv_ember.h5")
    print("Modified training data stored in: {}".format(train_set))
    print("Modified testing data stored in: {}".format(test_set))

    return agent_score



if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        usage='h4rm0ny',
    )
    
    parser.add_argument(
        '--agent',
        choices=['acer', 'ddqn', 'random'],
        default='acer',
        help='The agent.'
    )
    parser.add_argument(
        '--dynamic',
        action='store_true', 
        help='Uptrain the model.',
    )
    parser.add_argument(
        '--rounds', 
        type=int, 
        default=10000,
        help='an integer for the rounds.')
    parser.add_argument(
        '--gpuid', 
        type=int, 
        default=0,
        help='an integer for the gpuid.')
    parser.add_argument(
        '--train',
        action='store_true', 
        help='Uptrain the model.',
    )
    
    flags = parser.parse_args()
    print(flags)
    research(agent=flags.agent, dynamic=flags.dynamic, rounds=flags.rounds, gpu_id=flags.gpuid, train=flags.train)