#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import Image, LaserScan
import cv2
from cv_bridge import CvBridge

from datetime import datetime
from pathlib import Path
import pickle

DO_LOGGING = False # save image and lidar to ../logging/

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
log_path = ROOT / 'logging'

def publish_images():

    # launch parameters
    try:
        CAMERA_WIDTH = int(rospy.get_param('camera_width'))
        CAMERA_HEIGHT = int(rospy.get_param('camera_height'))
        # TODO pridat cestu k yaml filu s geometrii konfigurace lidar-kamera
    except Exception:
        print("Can't load Camera sizes from ROS parameter server, using defaults.")
        CAMERA_WIDTH = 640
        CAMERA_HEIGHT = 360

    # node setup
    rospy.loginfo('Initializing webcamera stream.')
    rospy.init_node('webcam_stream', anonymous=False)
    pub = rospy.Publisher('video_frames', Image, queue_size=1)
    rate = rospy.Rate(10) # 10 Hz
        
    # node funcionality
    width, height = CAMERA_WIDTH, CAMERA_HEIGHT # 1280 x 800, 640 x 360, 424 x 240
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
    br = CvBridge()

    i = 0 # counter for logging

    # webcam stream
    while not rospy.is_shutdown():

        i += 1

        ret, frame = cap.read()

        now = datetime.now()
        image_stamp = now.strftime("%Y_%m_%d_%H_%M_%S_%f")

        if ret == True:
            pub.publish(br.cv2_to_imgmsg(frame))

            # log image and lidar data
            if DO_LOGGING and i >= 5:
                
                lidar_msg = rospy.wait_for_message("/scan", LaserScan)
                
                now = datetime.now()
                lidar_stamp = now.strftime("%Y_%m_%d_%H_%M_%S_%f")

                # save original image
                cv2.imwrite(str(log_path / f'{image_stamp}_img_original.jpg'), frame,
                    [cv2.IMWRITE_JPEG_QUALITY, 90] )

                # save lidar beams
                with open(log_path / f'{image_stamp}_lidar.pkl', 'wb') as f:
                    pickle.dump(lidar_msg, f)

                # save time stamps
                # 1) after taking image, 2) after taking lidar
                with open(log_path / f'{image_stamp}_time_stamps.pkl', 'wb') as f:
                    pickle.dump([image_stamp, lidar_stamp], f)
                
                i = 0


        rate.sleep()


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