# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#


import glob
import sys
import yaml
from train_model.dataset_utils import prepare_eval_data_set
from train_model.helper import build_model, run_model
import torch
from torch.utils.data import DataLoader

CONFIG = "config.yaml"
MODELNAME = "best_model.pth"

if __name__ == '__main__':
    if len(sys.argv) < 2:
        exit("USAGE: python tools/eval_ensemble_on_val.py parent_dir \
             [ensemble sizes]")

    esbl_sizes = [int(a) for a in sys.argv[2:]]

    parent_dir = sys.argv[1]

    model_pths = [file for file in glob.glob(
        parent_dir + "/**/" + MODELNAME, recursive=True)]
    config_files = [c.replace(MODELNAME, CONFIG) for c in model_pths]

    if len(esbl_sizes) == 0:
        esbl_sizes = range(1, len(config_files)+1)

    config_file = config_files[0]

    with open(config_file, 'r') as f:
        config = yaml.load(f)

    batch_size = config['data']['batch_size']
    data_set_test = prepare_eval_data_set(**config['data'],
                                          **config['model'],
                                          verbose=True)
    data_reader_test = DataLoader(data_set_test,
                                  shuffle=False,
                                  batch_size=batch_size,
                                  num_workers=5)
    ans_dic = data_set_test.answer_dict

    accumulated_softmax = None
    final_result = {}
    n_model = 0
    for c_file, model_file in zip(config_files, model_pths):
        with open(c_file, 'r') as f:
            config = yaml.load(f)

        myModel = build_model(config, data_set_test)
        myModel.load_state_dict(torch.load(model_file)['state_dict'])

        question_ids, soft_max_result = run_model(myModel,
                                                  data_reader_test,
                                                  ans_dic.UNK_idx)

        if n_model == 0:
            final_result = soft_max_result