import os
import sys
import argparse
import string
import subprocess

from rlpytorch.utils import Logger


def eval_folder(folder):
    files = os.listdir(folder)
    models = [os.path.join(folder, f) for f in files if f.endswith('pt')]
    models = [os.path.abspath(m) for m in models]

    processes = []
    logs = []
    for _, model in enumerate(models):
        save_dir = model + '_replay'
        log_name = model + '_eval.log'
        log = open(log_name, 'wb')
        cmd = [
            './eval_executor.sh',
            model,
            save_dir,
        ]
        cmd = ' '.join(cmd)
        p = subprocess.Popen(
            [cmd],
            shell=True,
            cwd='../../../rts_scripts',
            stdout=log,
            stderr=subprocess.PIPE,
        )
        processes.append(p)
        logs.append((log_name, log))
        # break

    for i, p in enumerate(processes):
        p.wait()
        log_name, log = logs[i]
        log.close()
        log = open(log_name, 'r')
        result = log.readlines()[-1].strip()
        model_sig = models[i].split('/')[-1]
        print('%s: %s' % (model_sig, result))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--folder', type=str)
    parser.add_argument('--log_file', type=str, default='eval.log')
    args = parser.parse_args()

    logger_path = os.path.join(args.folder, args.log_file)
    print('logging to ', logger_path)
    sys.stdout = Logger(logger_path)
    eval_folder(args.folder)
