#!/usr/bin/env python

from __future__ import print_function
import time
import sys
import cv2
import numpy as np
import rospy

# Needed for publishing the messages
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

class VideoPublish:

    def __init__(self):
        self.frames = []
        self.max_frames = 100
        self.frame_index = 0

        rospy.init_node('video_publish')

        start_frame = rospy.get_param('~start_frame', 1)
        fps = rospy.get_param('~fps', 15)
        rate = rospy.Rate(fps)
        video_source = rospy.get_param('~video_source')
        frame_id = rospy.get_param('~frame_id', 'camera')
        height = rospy.get_param('~height', -1)
        width = rospy.get_param('~width', -1)

        single_frame = rospy.get_param('~single_frame', False)
        preview = single_frame or rospy.get_param('~preview', False)

        print('single_frame:', single_frame)
        print('preview:', preview)
        pub = rospy.Publisher('/color/image_raw', Image, queue_size=1000)
        cap = cv2.VideoCapture(video_source)
        if height > 0 and width > 0:
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
            cap.set(cv2.CAP_PROP_FPS, fps)
        bridge = CvBridge()

        frame_number = 0
        while frame_number+1 < start_frame:
            ret, im = self.read_frame(cap)
            if not ret:
                break
            frame_number += 1

        while not rospy.is_shutdown():
            ret, im = self.read_frame(cap)
            if not ret:
                break
            frame_number += 1

            msg = bridge.cv2_to_imgmsg(im, 'bgr8')
            msg.header.stamp = rospy.Time.now()
            msg.header.frame_id = frame_id
            pub.publish(msg)

            if preview or single_frame:
                im_copy = im.copy()
                msg = 'Frame: {0:d}'.format(frame_number)
                cv2.putText(im_copy, msg, (10,40), cv2.FONT_HERSHEY_SIMPLEX,
                            1, (255,0,0), 2, cv2.LINE_AA)
                cv2.imshow('image', im_copy)

            if not single_frame:
                if preview:
                    key = cv2.waitKey(1)
                    if key == 27:
                        break
                rate.sleep()
            else:
                key = cv2.waitKey(0)
                if key == 27:
                    break
                elif key == ord(';'):
                    frame_number -= self.backup_frames(2)
                elif key == ord(']'):
                    frame_number += self.skip_frames(cap, 9)
                elif key == ord('['):
                    frame_number -= self.backup_frames(11)
                elif key == ord('='):
                    frame_number += self.skip_frames(cap, 99)
                elif key == ord('/'):
                    single_frame = False
                    preview = True

    def skip_frames(self, cap, n):
        skipped = 0
        for i in range(n):
            ret, im = self.read_frame(cap)
            if not ret:
                break
            skipped += 1
        return skipped

    def backup_frames(self, n):
        n = min(self.frame_index, n)
        self.frame_index -= n
        return n

    def read_frame(self, cap):
        if self.frame_index < len(self.frames):
            frame = self.frames[self.frame_index]
            self.frame_index += 1
            return True, frame
        else:
            ret, frame = cap.read()
            if ret:
                self.frames.append(frame)
                if len(self.frames) > self.max_frames:
                    self.frames.pop(0)
                else:
                    self.frame_index += 1
            return ret, frame

if __name__ == '__main__':
    try:
        VideoPublish()
    except rospy.ROSInterruptException:
        pass
