#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import argparse
from loguru import logger
import multiprocessing as mp

import megengine as mge
import megengine.distributed as dist

from yolox.core import Trainer
from yolox.exp import get_exp
from yolox.utils import configure_nccl


def make_parser():
    parser = argparse.ArgumentParser("YOLOX train parser")
    parser.add_argument("-expn", "--experiment-name", type=str, default=None)
    parser.add_argument("-n", "--name", type=str, default=None, help="model name")

    parser.add_argument("-b", "--batch-size", type=int, default=64, help="batch size")
    parser.add_argument(
        "-d", "--devices", default=None, type=int, help="device for training"
    )
    parser.add_argument(
        "-f",
        "--exp_file",
        default=None,
        type=str,
        help="plz input your expriment description file",
    )
    parser.add_argument(
        "--resume", default=False, action="store_true", help="resume training"
    )
    parser.add_argument("-c", "--ckpt", default=None, type=str, help="checkpoint file")
    parser.add_argument(
        "--num_machine", default=1, type=int, help="num of node for training"
    )
    parser.add_argument(
        "--machine_rank", default=0, type=int, help="node rank for multi-node training"
    )
    parser.add_argument(
        "--sync_level", type=int, default=None, help="config sync level, use 0 to debug"
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    return parser


@logger.catch
def main(exp, args):
    if not args.experiment_name:
        args.experiment_name = exp.exp_name

    # set environment variables for distributed training
    configure_nccl()

    # enable dtr to avoid CUDA OOM
    mge.dtr.enable()

    if args.sync_level is not None:
        # NOTE: use sync_level = 0 to debug mge error
        from megengine.core._imperative_rt.core2 import config_async_level
        logger.info("Using aysnc_level {}".format(args.sync_level))
        
config_async_level(args.sync_level)