import argparse

from tracking.data import split_and_pad
from tracking.tracking import track
from tracking.visualisation import visualise_tracking


def create_arg_parser():
    parser = argparse.ArgumentParser(description='Track Objects in a Video.')

    subparsers = parser.add_subparsers(title='sub-commands',
                                       description='Valid sub-commands.')

    # Track
    parser_track = subparsers.add_parser('track', help='Track Objects in a Video.')
    parser_track.add_argument('pose_data', type=str, help='.json file containing the keypoints and confidence scores '
                                                          'of the detected objects in a video. .json file must be in '
                                                          'the COCO format.')
    parser_track.add_argument('frames', type=str, help='Directory containing the frames of the video.')
    parser_track.add_argument('--use_bounding_box_strategy', action='store_true',
                              help='Use the bounding box strategy instead of sparse optical flow.')
    parser_track.add_argument('--video_resolution', default='856x480', type=str,
                              help='Resolution of the video from where the objects were detected.')
    parser_track.add_argument('--detection_quantile', default=10, type=int,
                              help='Threshold for removal of poor detections. Any detection with a detection score '
                                   'below the specified quantile is removed.')
    parser_track.add_argument('--keypoint_confidence_threshold', default=0.30, type=float,
                              help='Threshold for removal of poor keypoints. Any keypoint with a detection confidence '
                                   'below this threshold is considered to be a missing keypoint. The detection '
                                   'confidence of a keypoint can range from 0 to 1.')
    parser_track.add_argument('--minimum_keypoints', default=7, type=int,
                              help='Any detection with less than the specified minimum number of keypoints is '
                                   'removed. This happens after poor keypoints are considered missing.')
    parser_track.add_argument('--look_back', default=5, type=int,
                              help='How many frames to look back while searching for a matching object.')
    parser_track.add_argument('--bounding_box_edge_ratio', default=0.2, type=float,
                              help='Create a small bounding box around each keypoint, where the length of the edges '
                                   'of this bounding box are a percentage of the length of the edges of the detected '
                                   'object\'s bounding box. These small bounding boxes are utilised by the optical '
                                   'flow to measure how close the pose predicted by optical flow is to the pose '
                                   'detected by the object detector. On the other hand, the bounding box strategy '
                                   'utilises them to compute their IOU.')
    parser_track.add_argument('--num_top_keypoints', default=7, type=int,
                              help='The number of keypoints to consider when computing the average IOU between the '
                                   'small bounding boxes created around each keypoint. Only meaningful if using the '
                                   'bounding box strategy.')
    parser_track.add_argument('--matching_threshold', default=0.0, type=float,
                              help='Threshold for tracking confirmation. If the matching score between two objects in '
                              'distinct frames is above this threshold, then these two objects are considered to be '
                              'the same. The matching score between two objects can range from 0 to 1.')
    parser_track.add_argument('--write_tracking', type=str,
                              help='File to write tracking results. Tracking results will be saved in .csv format, '
                                   'where the first column contains the frame number, the second column contains '
                                   'the object id and all other columns contain the keypoints of the tracked objects.')

    parser_track.set_defaults(func=track)

    # Split and Pad Tracking
    parser_split_and_pad = subparsers.add_parser('split_and_pad',
                                                 help='Split Tracked Trajectories by Object and Pad Missing Frames.')
    parser_split_and_pad.add_argument('tracked_trajectories', type=str,
                                      help='.csv file containing tracked trajectories of objects in a video.')
    parser_split_and_pad.add_argument('--write_split_trajectories', type=str,
                                      help='Directory to write results of the splitting. For each object tracked, a '
                                           '.csv file is created containing only this object\'s trajectory.')

    parser_split_and_pad.set_defaults(func=split_and_pad)

    # Visualise Tracking
    parser_visualise_tracking = subparsers.add_parser('visualise_tracking',
                                                      help='Visualise Tracked Trajectories.')
    parser_visualise_tracking.add_argument('tracked_trajectories', type=str,
                                           help='.csv file containing tracked trajectories of objects in a video.')
    parser_visualise_tracking.add_argument('frames', type=str, help='Directory containing the frames of the video.')
    parser_visualise_tracking.add_argument('write_rendered_tracking', type=str,
                                           help='Directory to write rendered frames.')
    parser_visualise_tracking.add_argument('--draw_only_bounding_boxes', action='store_true',
                                           help='Draw only the bounding boxes of the objects.')

    parser_visualise_tracking.set_defaults(func=visualise_tracking)

    return parser


def main():
    arg_parser = create_arg_parser()
    args = arg_parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()
