#!/usr/bin/env python
import ecto

from ecto_opencv.highgui import imshow
from ecto_opencv.imgproc import cvtColor, Conversion
from ecto.opts import scheduler_options, run_plasm, cell_options
from ecto_image_pipeline.io.source import create_source
from ecto_image_pipeline.io.source import add_camera_group

from object_recognition_capture.orb_capture import OrbPoseEstimator

if __name__ == '__main__':
    def parse_args():
        import argparse
        parser = argparse.ArgumentParser(description='Estimate the pose of an ORB template.')

        scheduler_options(parser.add_argument_group('Scheduler'))
        factory = cell_options(parser, OrbPoseEstimator, 'track')
        add_camera_group(parser)

        options = parser.parse_args()
        options.orb_factory = factory
        return options

    options = parse_args()
    plasm = ecto.Plasm()

    #setup the input source, grayscale conversion
    source = create_source('image_pipeline','OpenNISource',outputs_list=['K_depth', 'K_image', 'image', 'depth', 'points3d', 'mask_depth'],res=options.res,fps=options.fps)
    rgb2gray = cvtColor('Grayscale', flag=Conversion.RGB2GRAY)
    plasm.connect(source['image'] >> rgb2gray ['image'])

    pose_est = options.orb_factory(options, 'ORB Tracker')

    #convenience variable for the grayscale
    img_src = rgb2gray['image']

    #connect up the pose_est
    plasm.connect(img_src >> pose_est['image'],
                  source['image','mask_depth','K_image','points3d'] >> pose_est['color_image','mask','K_image','points3d'],
                  )

    display = imshow('orb display', name='Pose')
    plasm.connect(pose_est['debug_image'] >> display['image'],
                  )

    run_plasm(options, plasm, locals=vars())
