#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Software License Agreement (BSD License)
#
#  Copyright (c) 2015, Willow Garage, Inc.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#   * Redistributions in binary form must reproduce the above
#     copyright notice, this list of conditions and the following
#     disclaimer in the documentation and/or other materials provided
#     with the distribution.
#   * Neither the name of the Willow Garage nor the names of its
#     contributors may be used to endorse or promote products derived
#     from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
#  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
#  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
#  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
#  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
#  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
"""Save images of multiple topics with timestamp synchronization.

Usage: ros2 run image_view extract_images_sync --ros-args -p inputs:='[<topic_0>, <topic_1>]'
"""

import sys
import cv2
import cv_bridge
import rclpy

from rclpy.node import Node
from rclpy.duration import Duration
from sensor_msgs.msg import Image
from message_filters import ApproximateTimeSynchronizer, TimeSynchronizer, Subscriber

class ExtractImagesSync(Node):

    def __init__(self):
        super().__init__('extract_images_sync')
        self.get_logger().info('Extract_Images_Sync Node has been started')
        self.seq = 0
        self.last_frame_time = self.get_clock().now()

        self.fname_fmt = self.declare_parameter(
            'filename_format', 'frame%04i_%i.jpg').value

        # Do not scale dynamically by default.
        self.do_dynamic_scaling = self.declare_parameter(
            'do_dynamic_scaling', False).value

        # Limit the throughput to 10fps by default.
        self.sec_per_frame = self.declare_parameter(
            'sec_per_frame', 0.0).value

        img_topics = self.declare_parameter('inputs', None).value

        if img_topics is None:
            self.get_logger().warn("""\
extract_images_sync: Parameter 'inputs' has not been specified! \
Typical command-line usage:
\t$ ros2 run image_view extract_images_sync --ros-args -p inputs:=<image_topic>
\t$ ros2 run image_view extract_images_sync --ros-args -p inputs:='[<image_topic>, <image_topic>]'""")
            sys.exit(1)

        if not isinstance(img_topics, list):
            img_topics = [img_topics]

        subs = []
        for t in img_topics:
            subs.append(Subscriber(self, Image, t))

        self.approximate_sync = self.declare_parameter('approximate_sync', False).value;
        if self.approximate_sync:
            sync = ApproximateTimeSynchronizer(subs, 100, slop=0.1)
        else:
            sync = TimeSynchronizer(subs, 100)
        sync.registerCallback(self.save)

    def save(self, *imgmsgs):
        delay = self.get_clock().now() - self.last_frame_time

        # Decimation is enabled if it is set to greater than
        # zero. If it is enabled, exit if delay is less than
        # the decimation rate specified.
        if (self.sec_per_frame > 0) and (delay < Duration(seconds=self.sec_per_frame)):
            return
        
        # Update the current time
        self.last_frame_time = self.get_clock().now()

        seq = self.seq
        bridge = cv_bridge.CvBridge()
        for i, imgmsg in enumerate(imgmsgs):
            img = bridge.imgmsg_to_cv2(imgmsg)
            # Default to the encoding in the image message
            encoding_in = imgmsg.encoding
            if not encoding_in:
                # Encoding is not specified, try to automatically determine
                channels = img.shape[2] if img.ndim == 3 else 1
                encoding_in = bridge.dtype_with_channels_to_cvtype2(
                    img.dtype, channels)
            img = cv_bridge.cvtColorForDisplay(
                img, encoding_in=encoding_in, encoding_out='',
                do_dynamic_scaling=self.do_dynamic_scaling)
            fname = self.fname_fmt % (seq, i)
            print('Save image as {0}'.format(fname))
            cv2.imwrite(fname, img)
        self.seq = seq + 1

def main(args=None):
    rclpy.init(args=args)
    extractor = ExtractImagesSync()
    rclpy.spin(extractor)
    rclpy.shutdown()

if __name__ == '__main__':
    main()
