import pdb
import numpy as np
import torch.utils.data as data
import utils
from options import *
from config import *
from test import *
from model import *
from tensorboard_logger import Logger
from thumos_features import *


if __name__ == "__main__":
    args = parse_args()
    if args.debug:
        pdb.set_trace()

    if os.path.isfile(args.config):
        cfg = load_config(args.config)
    else:
        raise ValueError("Config file does not exist.")
    print(cfg)

    cfg = init_args(cfg)
    config = Config(cfg)

    worker_init_fn = None

    if config.seed >= 0:
        utils.set_seed(config.seed)
        worker_init_fn = np.random.seed(config.seed)

    utils.save_config(config, os.path.join(config.output_path, "config.txt"))

    net = Model(config.len_feature, config.num_classes, config.r_act)
    net = net.cuda()
    dataset_name = config.dataset_name

    test_loader = data.DataLoader(
        build_dataset(dataset_name=dataset_name, data_path=config.data_path, mode='test',
                      modal=config.modal, feature_fps=config.feature_fps,
                      num_segments=-1, sampling='random',
                      supervision='point', seed=config.seed),
        batch_size=1,
        shuffle=False, num_workers=config.num_workers,
        worker_init_fn=worker_init_fn)

    test_info = utils.build_test_info(dataset_name)
    
    logger = Logger(config.log_path)

    net.load_state_dict(torch.load(config.model_file))

    test(net, config, logger, test_loader, test_info, 0)

    if 'thumos14' in dataset_name:
        utils.save_best_record_thumos(test_info, os.path.join(config.output_path, "best_record.txt"))
    elif 'activitynet13' in dataset_name:
        utils.save_best_record_activity(test_info, os.path.join(config.output_path, "best_record.txt"))
