import os
import math
import argparse
from multiprocessing import Process, Queue

from tqdm import tqdm
import numpy as np
import megengine as mge
from megengine import jit

from config import config
import network
import dataset
import misc_utils

if_set_nms = True

def eval_all(args):
    # model_path
    saveDir = config.model_dir
    evalDir = config.eval_dir
    misc_utils.ensure_dir(evalDir)
    model_file = os.path.join(saveDir, 
            'epoch_{}.pkl'.format(args.resume_weights))
    assert os.path.exists(model_file)
    # load data
    records = misc_utils.load_json_lines(config.eval_source)
    # multiprocessing
    num_records = len(records)
    num_devs = args.devices
    num_image = math.ceil(num_records / num_devs)
    result_queue = Queue(1000)
    procs = []
    all_results = []
    for i in range(num_devs):
        start = i * num_image
        end = min(start + num_image, num_records)
        split_records = records[start:end]
        proc = Process(target=inference, args=(
                model_file, i, split_records, result_queue))
        proc.start()
        procs.append(proc)
    pbar = tqdm(total=num_records, ncols=50)
    for i in range(num_records):
        t = result_queue.get()
        all_results.append(t)
        pbar.update(1)
    for p in procs:
        p.join()
    fpath = os.path.join(evalDir, 'dump-{}.json'.format(args.resume_weights))
    misc_utils.save_json_lines(all_results, fpath)

def inference(model_file, device, records, result_queue):
    @jit.trace(symbolic=False)
    def val_func():
        pred_boxes = net(net.inputs)
        return pred_boxes
    net = network.Network()
    net.eval()
    check_point = 
mge.load(model_file)