import os
from parlai.scripts.display_model import DisplayModel
from parlai.scripts.train_model import TrainModel

from GenericTools.StayOrganizedTools.utils import timeStructured
from ariel_tests.models.dialogue.parlai_agent import *
from ariel_tests.models.dialogue.hello_agent import *

CDIR = os.path.dirname(os.path.realpath(__file__))
EXPERIMENTS = os.path.join(CDIR, 'experiments')
if not os.path.isdir(EXPERIMENTS): os.mkdir(EXPERIMENTS)

time_string = timeStructured()
EXPERIMENT = os.path.join(EXPERIMENTS, "{}_td".format(time_string))
if not os.path.isdir(EXPERIMENT): os.mkdir(EXPERIMENT)



print('=' * 30)
DisplayModel.main(
    task='light_dialog',
    # model_file='models:light/biranker_dialogue/model',
    model='hello',
    num_examples=5,
    datatype='test',
)

print('=' * 30)
TrainModel.main(
    model='my_first_lstm',
    model_file=os.path.join(EXPERIMENT, 'model'),
    task='empathetic_dialogues',
    batchsize=64,
    validation_every_n_secs=10,
    max_train_time=2,
    num_epochs=2,
    verbose=True
)

