# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import os
import argparse
import datetime
import multiprocessing as mp
import time

import megengine as mge
import megengine.amp as amp
import megengine.autodiff as autodiff
import megengine.distributed as dist
import megengine.functional as F
import megengine.jit as jit
import megengine.optimizer as optim
from basecore.utils import log_every_n_seconds
from loguru import logger

from basecls.data.fake_data import FakeDataLoader
from basecls.layers import Preprocess
from basecls.utils import registers, set_nccl_env, set_num_threads


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", default="resnet50", type=str)
    parser.add_argument("--mode", default="eval", type=str)
    parser.add_argument("-d", "--device", default="gpu", type=str)
    parser.add_argument("--amp", default=0, type=int)
    parser.add_argument("--fastrun", action="store_true")
    parser.add_argument("--trace", action="store_true")
    parser.add_argument("--dtr", action="store_true")
    parser.add_argument("-b", "--batch-size", default=32, type=int)
    parser.add_argument("--channel", default=3, type=int)
    parser.add_argument("--height", default=224, type=int)
    parser.add_argument("--width", default=224, type=int)
    parser.add_argument("-n", "--world-size", default=8, type=int)
    parser.add_argument("--warm-iters", default=50, type=int)
    parser.add_argument("-t", "--total-iters", default=200, type=int)
    parser.add_argument("--log-seconds", default=2, type=int)
    args = parser.parse_args()

    mp.set_start_method("spawn")

    set_nccl_env()
    set_num_threads()

    if args.world_size == 1:
        worker(args)
    else:
        
dist.launcher(worker, n_gpus=args.world_size)