"""Configurate arguments."""
import argparse

# 0: confidence, 1: point_shape, 2: offset_x, 3: offset_y, 4: cos(direction),
# 5: sin(direction)
NUM_FEATURE_MAP_CHANNEL = 6

def add_common_arguments(parser):
    """Add common arguments for training and inference."""
    parser.add_argument('--detector_weights', default=r'E:\pycharmProject\DMPR-PS\weights\dp_detector_299.pth',
                        help="The weights of pretrained detector.")
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml',
                        help='model.yaml path')
    parser.add_argument('--hyp', type=str, default='conf/hyp.scratch.yaml',
                        help='hyperparameters path')
    parser.add_argument('--depth_factor', type=int, default=32,
                        help="Depth factor.")
    parser.add_argument('--disable_cuda', action='store_true',
                        help="Disable CUDA.")
    parser.add_argument('--gpu_id', type=int, default=0,
                        help="Select which gpu to use.")


def get_parser_for_inference():
    """Return argument parser for inference."""
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', default='image', choices=['image', 'video'],
                        help="Inference image or video.")
    parser.add_argument('--video',
                        help="Video path if you choose to inference video.")
    parser.add_argument('--inference_slot', action='store_true', default=False,
                        help="Perform slot inference.")
    parser.add_argument('--thresh', type=float, default=0.3,
                        help="Detection threshold.")
    parser.add_argument('--save', action='store_true', default=True,
                        help="Save detection result to file.")
    parser.add_argument('--yoloimg-size', type=int, default=640,
                        help='inference size (pixels)')
    parser.add_argument('--dmprimg-size', type=int, default=640,
                        help='inference size (pixels)')
    parser.add_argument('--augment', action='store_true',
                        help='augmented inference')
    parser.add_argument('--conf-thres', type=float, default=0.25,
                        help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45,
                        help='IOU threshold for NMS')
    parser.add_argument('--scale-ratio', type=float, default=0.5,
                        help='detected box scale ratio')
    parser.add_argument('--border', type=int, default=80,
                        help='The valid border to boundary')
    parser.add_argument('--ovlap-thres', type=float, default=0.6, help='overlap threshold for OBS')
    parser.add_argument('--agnostic-nms', action='store_true',
                        help='class-agnostic NMS')
    parser.add_argument('--classes', nargs='+', type=int,
                        help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--dmpr-thresh', type=float, default=0.3,
                        help="Detection threshold.")
    # STDC
    parser.add_argument('--n-classes', type=int, default=2, help='number of classes for segment')
    parser.add_argument('--backbone', type=str, default='STDCNet813', help='STDC backbone')
    parser.add_argument('--respth', type=str, default='weights/urbanManagement/STDC/model_final_1023.pth',
                        help='The weights of STDC')
    parser.add_argument('--stdc-new-hw', nargs='+', type=int, default=[360, 640], help='The new hw of STDC')
    parser.add_argument('--use-boundary-2', type=bool, default=False, help='')
    parser.add_argument('--use-boundary-4', type=bool, default=False, help='')
    parser.add_argument('--use-boundary-8', type=bool, default=False, help='')
    parser.add_argument('--use-boundary-16', type=bool, default=False, help='')
    parser.add_argument('--use-conv-last', type=bool, default=False, help='')

    add_common_arguments(parser)
    return parser