# parser.py
# Vojtěch Golda
import argparse

parser = argparse.ArgumentParser(
                    prog='python run.py',
                    description='This program detects trams.',
                    epilog='')

parser.add_argument('-p', '--model_path', help="Model checkpoint path.", type=str, required=True)
parser.add_argument('-g', action='store_true', help="Use CUDA/ROCm.")

parser.add_argument('-i', '--input_file', help="Input file path.", type=str)
parser.add_argument('-o', '--output_file', help="Output file path.", type=str)

parser.add_argument('-f', '--camera', help="Webcamera index. Overrides -i.", type=int)
parser.add_argument('-n', '--nth', help="Only process every nth frame. Only affects detections. Default 1.", type=int, default=1)
parser.add_argument('-s', action='store_true', help="Output detected tram coordinates to standard output.")
parser.set_defaults(s=False)
parser.add_argument('-w', action='store_true', help="Output detected trams to a window.")
parser.set_defaults(w=False)

parser.add_argument('-t', action='store_true', help="Switch to training mode.")
parser.set_defaults(t=False)
parser.add_argument('-a', '--train_dataset', help="Training dataset path. Mandatory if -t is used.", type=str)
parser.add_argument('-b', '--validation_dataset', help="Validation dataset path. Mandatory if -t and validation is used.", type=str)
parser.add_argument('-l', '--learning_rate', help="Learning rate. Default 0.0001.", type=float, default=0.0001)
parser.add_argument('-x', '--batch_size', help="Batch size. Default 5.", type=int, default=5)
parser.add_argument('-e', '--epoch_count', help="Epoch count. Default infinite.", type=int)


parser.add_argument('-v', action='store_true', help="Perform validation after each epoch.")
parser.set_defaults(t=False)


parser.add_argument('-m', '--model_name', help="Name of the model to use.", choices=[
        "FCOS", "Faster_R_CNN_MobileNet", "Faster_R_CNN_ResNet", "RetinaNet", "SSD", "SSDLite", "TramNet"
], required=True)

args = parser.parse_args()

