"""
    train.py
        This class specifies the training flow of the agent. From here the alternating exploration and determination of causal objects is performed.
"""

import argparse
import time
import random

from runnables import RunnableMiniHack, RunnableUnlockPickup

import numpy as np
import pickle


"""
    Main entry point.
"""
runner = None
trainer = None
num_views = 5
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # parser.add_argument('--n_train', type=int, default=1)
    # parser.add_argument('--n_episodes', type=int, default=1000)
    parser.add_argument('--env_name', type=str, default=None)
    parser.add_argument('--file_name', type=str, default='cycle')
    parser.add_argument('--run_num', type=int, default=None)
    parser.add_argument('--num_steps', type=int, default=None)
    parser.add_argument('--epsilon', type=float, default=None)
    parser.add_argument('--gamma', type=float, default=None)
    parser.add_argument('--alpha', type=float, default=None)
    parser.add_argument('--rho', type=float, default=1.0)


    args = parser.parse_args()

    file_name = args.file_name

    save = True

    if args.env_name is not None:
        print(args.env_name.split('-'))
        if args.env_name.split('-')[0] == 'MiniGrid':
            num_views = 5
            runnable = RunnableUnlockPickup(num_views)
        elif args.env_name.split('-')[0] == 'MiniHack':
            num_views = 4
            runnable = RunnableMiniHack(num_views)
    else:
        runnable = RunnableUnlockPickup(num_views)
    
    runnable.run(args, save)

