import argparse
import os
import sys

import horovod.torch as hvd
import torch

if __name__ == "__main__" and __package__ is None:
    sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
    __package__ = "BaseDetector"

from .config import get_cfg_defaults
from .utils.gpu_utils import set_gpu
from .engine import BaseDetection3D

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="full functional execute script of Detection module.")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-t", "--train", help="model training", action="store_true")
    group.add_argument("-i", "--inference", help="model inference", action="store_true")
    parser.add_argument("-c", "--config", type=str, default="config/config.yaml", help="config file path")
    args = parser.parse_args()

    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.config)
    cfg.freeze()
    print(cfg)

    # set gpu
    # os.environ['CUDA_VISIBLE_DEVICES'] = "1,2"
    set_gpu(cfg.TRAINING.NUM_GPUS, used_percent=0.1)

    # Horovod: init
    hvd.init()
    torch.cuda.set_device(hvd.local_rank())

    if hvd.rank() == 0:
        print(cfg)

    if args.train:
        detector = BaseDetection3D(cfg)
        detector.do_train()
    elif args.inference:
        detector = BaseDetection3D(cfg, "test")
        detector.do_test()