#!/usr/bin/env python2
print "importing libraries"
import rosbag
import cv_bridge

import sm
from sm import PlotCollection
import aslam_cv as acv
import aslam_cameras_april as acv_april
import aslam_cv_backend as acvb
import kalibr_common as kc
import kalibr_camera_calibration as kcc
import kalibr_rs_camera_calibration as rscc

import cv2
import os
import numpy as np
import pylab as pl
import argparse
import sys
import random
import signal

np.set_printoptions(suppress=True)


def __initBagDataset(bagfile, topic, from_to):
    print "\tDataset:          {0}".format(bagfile)
    print "\tTopic:            {0}".format(topic)
    reader = kc.BagImageDatasetReader(bagfile, topic, bag_from_to=from_to)
    print "\tNumber of images: {0}".format(reader.numImages())
    return reader


# available models
cameraModels = {
    # rolling shutter
    'pinhole-radtan-rs': acvb.DistortedPinholeRs,
    'pinhole-equi-rs': acvb.EquidistantPinholeRs,
    'omni-radtan-rs': acvb.DistortedOmniRs,
    # global shutter
    'pinhole-radtan': acvb.DistortedPinhole,
    'pinhole-equi': acvb.EquidistantPinhole,
    'omni-radtan': acvb.DistortedOmni
}


def __signal_exit(signal, frame):
    sm.logWarn("Shutdown requested! (CTRL+C)")
    sys.exit(2)


def __parseArgs():

    class KalibrArgParser(argparse.ArgumentParser):

        def error(self, message):
            self.print_help()
            sm.logError('%s' % message)
            sys.exit(2)

        def format_help(self):
            formatter = self._get_formatter()
            formatter.add_text(self.description)
            formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
            for action_group in self._action_groups:
                formatter.start_section(action_group.title)
                formatter.add_text(action_group.description)
                formatter.add_arguments(action_group._group_actions)
                formatter.end_section()
            formatter.add_text(self.epilog)
            return formatter.format_help()

    usage = """
    Example usage to calibrate a camera system with a single rolling shutter camera using an aprilgrid.

    cam: pinhole model with radial-tangential distortion and rolling shutter for a single camera

    %(prog)s --model pinhole-radtan-rs --target aprilgrid.yaml \\
              --bag MYROSBAG.bag --topic /cam0/image_raw \\
              --inverse-feature-variance 1 --frame-rate 20

    example aprilgrid.yaml:
        target_type: 'aprilgrid'
        tagCols: 6
        tagRows: 6
        tagSize: 0.088  #m
        tagSpacing: 0.3 #as a fraction of tagSize"""

    parser = KalibrArgParser(description='Calibrate the intrinsics of a single rolling shutter camera.', usage=usage)
    parser.add_argument(
        '--model',
        dest='model',
        help='The camera model to estimate. '
        'Currently supported models are {0}, where the suffix \'-rs\' identifies rolling shutter models.'.format(
            ', '.join(cameraModels.keys())),
        required=True)
    parser.add_argument('--frame-rate',
                        dest='framerate',
                        type=int,
                        help='Approximate framerate of the camera.',
                        required=True)

    groupSource = parser.add_argument_group('Data source')
    groupSource.add_argument('--bag', dest='bagfile', help='The bag file with the data.')
    groupSource.add_argument('--topic', dest='topic', help='The image topic.', required=True)
    groupSource.add_argument('--bag-from-to',
                             metavar='bag_from_to',
                             type=float,
                             nargs=2,
                             help='Use the bag data starting from up to this time [s].')

    groupTarget = parser.add_argument_group('Calibration target configuration')
    groupTarget.add_argument('--target',
                             dest='targetYaml',
                             help='Calibration target configuration as yaml file.',
                             required=True)
    groupTarget.add_argument('--feature-variance',
                             dest='featureVariance',
                             type=float,
                             help='Estimated variance of the feature detector.',
                             required=True)

    groupOpt = parser.add_argument_group('Optimization options')
    groupOpt.add_argument('--max-iter',
                          type=int,
                          default=30,
                          dest='max_iter',
                          help='Max. iterations (default: %(default)s).',
                          required=False)

    outputSettings = parser.add_argument_group('Output options')
    outputSettings.add_argument('--verbose',
                                action='store_true',
                                dest='verbose',
                                help='Enable (really) verbose output (disables plots).')
    outputSettings.add_argument('--show-extraction',
                                action='store_true',
                                dest='showextraction',
                                help='Show the calibration target extraction. (disables plots).')

    # print help if no argument is specified
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(2)

    # Parse the argument list
    try:
        parsed = parser.parse_args()
    except:
        sys.exit(2)

    # there is an issue with the gtk plot widget, so we cant plot if we have opencv windows open...
    # --> disable the plots in these special situations
    if parsed.showextraction or parsed.verbose:
        parsed.dontShowReport = True

    return parsed


# perform corner extraction from bagfile
# returns the observations
def __extractObservations(cameraGeometry, verbose=False):
    # extract the targets
    multithreading = False  # not verbose
    observations = kc.extractCornersFromDataset(
        cameraGeometry.dataset,
        cameraGeometry.ctarget.detector,
        multithreading=multithreading,
        # transformation estimation will fail with rs cameras and significant distortions
        noTransformation=True)

    return observations


def main():
    parsed = __parseArgs()

    # logging modes
    if parsed.verbose:
        sm.setLoggingLevel(sm.LoggingLevel.Debug)
    else:
        sm.setLoggingLevel(sm.LoggingLevel.Info)

    # register signal handler
    signal.signal(signal.SIGINT, __signal_exit)

    ######
    # load bagfile and extract targets:
    targetConfig = kc.CalibrationTargetParameters(parsed.targetYaml)
    dataset = __initBagDataset(parsed.bagfile, parsed.topic, parsed.bag_from_to)

    # create camera
    cameraModel = cameraModels[parsed.model]
    cameraGeometry = kcc.CameraGeometry(cameraModel,
                                        targetConfig,
                                        dataset,
                                        verbose=(parsed.verbose or parsed.showextraction))

    # extract observations
    observations = __extractObservations(cameraGeometry, (parsed.verbose or parsed.showextraction))

    # Calibration Configuration
    config = rscc.RsCalibratorConfiguration()
    config.inverseFeatureCovariance = 1.0 / parsed.featureVariance
    config.maxNumberOfIterations = parsed.max_iter
    config.framerate = parsed.framerate
    config.timeOffsetConstantSparsityPattern = 0.08
    config.timeOffsetPadding = 0.5
    config.targetSize = cameraGeometry.ctarget.grid.size()
    # camera parameters, here I fix the camera intrinsics, only optimize linedelay
    config.intrinsics = np.array([
        823.906785813718, 824.704794976756, 648.608565867332, 314.810494971340, -0.312098601430490, 0.0928470270344407,
        -1.93958495467811e-05, -0.000132104569851275
    ])
    # don't optimize the projection and distortion parameters
    config.estimateParameters['intrinsics'] = False
    config.estimateParameters['distortion'] = False
    config.adaptiveKnotPlacement = True

    ######
    # Calibrate
    calibrator = rscc.RsCalibrator()
    calibrator.calibrate(cameraGeometry, observations, config)


if __name__ == "__main__":
    try:
        main()
    except Exception, e:
        sm.logError("Exception: {0}".format(e))
        sys.exit(-1)
