import os,argparse
from utils import *
from train import *

##### CONFIG
config = {
    "exp_name":            "BasePopulation", # Experiment name
    "seed":                 0,             # seed
    "nb_agents":            2,             # Number of agents
    "nb_features":          5,
    "assoc_lr":             0.0001,        # LR for referent-utterance associations
    "action_lr":            0.01,          # LR for action generation
    "action_bs":            64,            # Nb of simultaneous search during action generation
    "action_it":            100,           # Nb of action generation iterations
    "action_size":          40,            # Size of agent's actions vectors
    "referents_bs":         32,            # Nb of games per training iteration
    "max_iterations":       10000,          # Max nb of training iterations
    "use_img_perspectives": True,          # Convert vectors referents into MNIST compositions?
    "ss_class":             "dmp",         # Sensorimotor system class
    "ss_params":            {"n_bfs":20, "dt":1e-1, "n":10, "d":52, "th":1e-2}, # Sensorimotor system params
    "shared_perspective":   False,
    "use_baseline":         True
}
################################################################################
parser = argparse.ArgumentParser(description='Training')
parser.add_argument('--nb_agents',   type=int, default=2)
parser.add_argument('--exp_name', type=str, default="Base")
parser.add_argument('--seed',     type=int, default=0)
parser.add_argument('--nb_features', type=int, default=5)
parser.add_argument('--assoc_lr',   type=float, default=0.0001)
parser.add_argument('--action_lr',  type=float, default=0.01)
parser.add_argument('--action_bs',   type=int, default=64)
parser.add_argument('--action_it',   type=int, default=100)
parser.add_argument('--action_size', type=int, default=40)
parser.add_argument('--ss_n',        type=int, default=10)
parser.add_argument('--referents_bs', type=int, default=32)
parser.add_argument('--max_iterations', type=int, default=10000)
parser.add_argument('--shared_perspective', action='store_true')
parser.add_argument('--no_perspective',     action='store_true')
parser.add_argument('--no_baseline',        action='store_true')

args = parser.parse_args()
################################################################################

if __name__ == "__main__":

    ''' Args '''
    config["seed"]     = args.seed
    config["exp_name"] = args.exp_name

    config["nb_agents"]   = args.nb_agents
    config["nb_features"] = args.nb_features

    config["assoc_lr"]   = args.assoc_lr
    config["action_lr"]  = args.action_lr

    config["action_bs"]         = args.action_bs
    config["action_it"]         = args.action_it
    config["action_size"]       = args.action_size
    config["ss_params"]["n"]    = args.ss_n

    config["referents_bs"]     = args.referents_bs
    config["max_iterations"]   = args.max_iterations

    config["shared_perspective"]    = args.shared_perspective
    config["use_img_perspectives"]  = not args.no_perspective
    config["use_baseline"]          = not args.no_baseline

    ''' Experiment path '''
    exp_path = path+config["exp_name"]+"/"
    if not os.path.exists(exp_path):
        os.mkdir(exp_path)

    ''' Launch Experiment '''
    launch_exp(config)
