#!/usr/bin/env python
import os, sys, argparse
import ecto, ecto_ros, ecto_ros.ecto_sensor_msgs as ecto_sensor_msgs, ecto_ros.ecto_geometry_msgs as ecto_geometry_msgs
from ecto_opencv import highgui, imgproc
from object_recognition_core.db import tools, models
from object_recognition_core.db.cells import ObservationInserter
from object_recognition_core.db.tools import args_to_db_params
import object_recognition_core.db.tools as dbtools

def connect_observation_calc_with_mask_pose(sync, commit, object_id, session_id, db_params, debug=False):
    plasm = ecto.Plasm()
    #need to go from ROS message types to opencv types.
    depth_ci = ecto_ros.CameraInfo2Cv('camera_info -> cv::Mat')
    image_ci = ecto_ros.CameraInfo2Cv('camera_info -> cv::Mat')
    image = ecto_ros.Image2Mat()
    depth = ecto_ros.Image2Mat()
    mask = ecto_ros.Image2Mat()
    pose = ecto_ros.PoseStamped2RT()

    #conversions
    plasm.connect(
      sync["image"] >> image["image"],
      sync["depth"] >> depth['image'],
      sync['image_ci'] >> image_ci[:],
      sync['depth_ci'] >> depth_ci[:],
      sync['mask'] >> mask[:],
      sync['pose'] >> pose[:]
      )

    rgb = imgproc.cvtColor('bgr -> rgb', flag=imgproc.Conversion.BGR2RGB)
    gray = imgproc.cvtColor('rgb -> gray', flag=imgproc.Conversion.RGB2GRAY)
    plasm.connect(image[:] >> (rgb[:], gray[:]),
                  )
    if debug:
        image_display = highgui.imshow('image display', name='image')
        mask_display = highgui.imshow('mask display', name='mask')
        plasm.connect(rgb[:] >> image_display[:])
        plasm.connect(mask[:] >> mask_display[:])

    if commit:
        db_inserter = ObservationInserter("db_inserter", object_id=object_id, session_id=session_id,
                                                  db_params=db_params)
        plasm.connect(depth[:] >> db_inserter['depth'],
                  pose['R', 'T'] >> db_inserter['R', 'T'],
                  mask[:] >> db_inserter['mask'],
                  rgb[:] >> db_inserter['image'],
                  image_ci['K'] >> db_inserter['K'],
                  )
    return plasm

def compute_for_bag_with_mask_pose(bag, obj, args):
    import couchdb
    couch = couchdb.Server(args.db_root)
    db = dbtools.init_object_databases(couch)
    objects = db
    existing = models.Object.by_object_name(objects, key=obj.object_name)
    store_new = True
    if len(existing) > 0:
        print "It appears that there are %d object(s) with the same name." % len(existing)
        for x in existing:
            print x
            print 'Use the object id above? [y,n]'
            use_it = raw_input("")
            if 'y' in use_it.lower():
                store_new = False
                obj = x
                break
    else:
        store_new = True
    if store_new:
        obj.store(objects)
        print "Stored new object with id:", obj.id

    ImageBagger = ecto_sensor_msgs.Bagger_Image
    CameraInfoBagger = ecto_sensor_msgs.Bagger_CameraInfo
    PoseBagger = ecto_geometry_msgs.Bagger_PoseStamped


    print "Loading bag with :", bag
    baggers = dict(image=ImageBagger(topic_name='/camera/rgb/image_color'),
                 depth=ImageBagger(topic_name='/camera/depth/image'),
                 mask=ImageBagger(topic_name='/camera/mask'),
                 pose=PoseBagger(topic_name='/camera/pose'),
                 image_ci=CameraInfoBagger(topic_name='/camera/rgb/camera_info'),
                 depth_ci=CameraInfoBagger(topic_name='/camera/depth/camera_info'),
                 )
    sync = ecto_ros.BagReader('Bag Reader',
                              baggers=baggers,
                              bag=bag,
                             )

    sessions = db
    session = models.Session()
    session.object_id = obj.id
    session.bag_id = bag
    if args.commit:
        session.store(sessions)
    print "running graph"
    plasm = connect_observation_calc_with_mask_pose(sync, args.commit,
                                     str(session.object_id),
                                     str(session.id), args_to_db_params(args),
                                     args.visualize)

    from ecto.opts import run_plasm
    run_plasm(args, plasm, locals=vars())
    return session
def parse_args():
    parser = argparse.ArgumentParser(description='''Uploads a bag, with an object description to the db.''',
                                      fromfile_prefix_chars='@')
    parser.add_argument('-i', '--input', metavar='BAG_FILE', dest='bag', type=str, default='',
                       help='A bag file to upload.')
    parser.add_argument('-n', '--object_name', metavar='OBJECT_NAME', dest='object_name', type=str, default='')
    parser.add_argument('-d', '--description', metavar='DESCRIPTION', dest='description', type=str, default='')
    parser.add_argument('-a', '--author', metavar='AUTHOR_NAME', dest='author_name', type=str, default='')
    parser.add_argument('-e', '--email', metavar='EMAIL_ADDRESS', dest='author_email', type=str, default='')
    parser.add_argument('tags', metavar='TAGS', type=str, nargs='+', help='Tags to add to object description.')
    parser.add_argument('--visualize', dest='visualize', action='store_const', const=True, default=False,
                        help='Turn on visualization')
    dbtools.add_db_arguments(parser)
    from ecto.opts import scheduler_options
    group = parser.add_argument_group('Scheduler Options')
    scheduler_options(group)

    args = parser.parse_args()
    if not args.bag:
        parser.print_help()
        sys.exit(1)
    if not os.path.exists(args.bag):
        sys.stderr.write('bag file does not exist')
        #raise RuntimeError('bag file does not exist')
        sys.exit(1)
    return args

if "__main__" == __name__:
    args = parse_args()
    bag = args.bag
    obj = models.Object(object_name=args.object_name,
                        description=args.description,
                        tags=args.tags,
                        author_name=args.author_name,
                        author_email=args.author_email,
                        )

    if args.commit:
        session = compute_for_bag_with_mask_pose(bag, obj, args)
        print 'Uploaded session with id:', session.id
    else:
        print 'Did not upload. Please pass --commit to actually upload the bag.'

