#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import print_function
from argus_gui import Undistorter, DistortionProfile
import argparse
from argus_gui import ArgusError
from six.moves import map

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Takes a video and uses argus_gui to undistort it.')
    parser.add_argument('ifile', type=str, help = 'Input movie. Movie may be in any format the OpenCV can read')
    parser.add_argument('--ofile', type=str, help = 'Location and name of movie to write to. Must end in .mp4', default = '')
    parser.add_argument('--mode', default = "None")
    parser.add_argument('--model', default = "None")
    parser.add_argument('--frameint', type=int, help = 'Maximum number of frames inbetween key frames in the newly written MP4', default = '25')
    parser.add_argument('--crf', type=int, help = 'Compression quality level. Must be an integer between 0 and 63, 0 being no compression (lossless) and 63 being maximal compression. Consult FFMPEGs documentation for more information', default='12')
    parser.add_argument('--write', help = 'Whether or not to write the video.', action = 'store_true')
    parser.add_argument('--display', help = 'Whether or not to display the undistortion as it happends.', action = 'store_true')
    parser.add_argument('--tmp', type=str, help = 'Temporary directory where frames are stored before writing. The directory will be deleted upon the scripts deletion. Be careful!', default = None)
    parser.add_argument('--crop', help = 'whether or not to crop the video to approximately the undistorted region.', action = 'store_true')
    parser.add_argument('--omni', type=str, help = 'String of omnidirectional distortion coefficients separated by commas. See the documenation for the correct formatting of this string', default = None)
    parser.add_argument('--coefficients', type=str, help = 'String of pinhole model distortion coefficients separated by commas. See the documenation for the correct formatting of this string', default = None)
    parser.add_argument('--cmei', action = 'store_true', help = 'Specifies that the coefficients passed our omnidirectional coefficients of CMeis model in OpenCV 3')
    parser.add_argument('--copy', action = 'store_true', help = 'Specifies that the inputted video should be copied.  Helps deal with bad headers, missing frames, etc.')
    args = parser.parse_args()
                        
    ifile = args.ifile
    ofile = args.ofile
    frameint = args.frameint
    crf = args.crf
    write = args.write
    display = args.display
    tmpName = args.tmp
    crop = args.crop
    
    if not args.omni is None:
        try:
            omni = list(map(float, args.omni.split(',')))
        except:
            raise ArgusError('distortion coefficients must all be valid floats')
    else:
        omni = None
    
    if not args.coefficients is None:
        try:
            coefficients = list(map(float, args.coefficients.split(',')))
        except:
            raise ArgusError('distortion coefficients must all be valid floats')
    else:
        coefficients = None

    if (args.mode != "None") and (args.model != "None"):
        dis = DistortionProfile()
        if '(Fisheye)' in args.mode:
            omni = dis.get_coefficients(args.model, args.mode)
        else:
            coefficients = dis.get_coefficients(args.model, args.mode)

            
    if (not omni is None) or (not coefficients is None):
        und = Undistorter(ifile, omni, coefficients, args.cmei, copy = args.copy)
        und.undistort_movie(ofile, frameint, crf, write, display, tmpName, crop)
    else:
        print('No distortion profile found or supplied.  If mode and model specified please make sure they are available in Argus.')
