from rosbag import Bag
from cv_bridge import CvBridge
from cv2 import imwrite
import numpy as np
import time
import cv2


# import libtiff


def convertDepthImage32FC1to16UC1(image_in, scale=1000.0):
    """
    The assumption is that the image_in is a float 32 bit np array, which contains the range values in meters.
    The image_out will be a 16 bit unsigned int  [0, 65536] which will store the range in millimeters.
    As a convetion, nans will be set to the maximum number available for uint16
    :param image_in: np array with dtype float32
    :param scale: convertion from meters to milimieter if 1000
    :return: np array with dtype uint16
    """

    if image_in.dtype == np.uint16:
        raise ValueError("Cannot convert float32 to uint16 because image is already uint16.")
        # return image_in

    # mask_nans = np.isnan(image_in)
    image_mm_float = image_in * scale  # convert meters to millimeters
    image_mm_float = np.round(image_mm_float)  # round the millimeters
    image_out = image_mm_float.astype(np.uint16)  # side effect: nans become 0s
    return image_out


# import cv2.cv as cv

cv_bridge = CvBridge()

bag = Bag('/home/daniela/bagfiles/hpe_stopped_far.bag')

# topics /camera_2/rgb/image_raw/compressed /camera_3/rgb/image_raw/compressed /camera_4/rgb/image_raw/compressed

# topic_rgb = '/camera_2/rgb/image_raw/compressed'
# window_name = 'image'
# for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
#     np_arr = np.frombuffer(msg.data, np.uint8)
#     img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
#     # img = cv_bridge.imgmsg_to_cv2(msg, "mono8")
#     p = '../images/camera_2/' + str(t) + '.png'
#     imwrite(p, img)

# topic_rgb = '/camera_1/rgb/image_color'
# window_name = 'image'
# i = 0
# for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
#     img = cv_bridge.imgmsg_to_cv2(msg)
#     # img=cv2.ConvertScaleAbs(img,256.0E)
#     p = '../images/with_depth/camera_1/' + str(i) + '.png'
#     # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
#     # print(img)
#     # print(img.dtype)
#     # cv2.imshow(window_name, img)
#     # cv2.waitKey(0)
#     imwrite(p, img)
#     i = i + 1

# topic_rgb = '/camera_4/rgb/image_raw/compressed'
# window_name = 'image'
# for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
#     np_arr = np.frombuffer(msg.data, np.uint8)
#     img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
#     # img = cv_bridge.imgmsg_to_cv2(msg, "mono8")
#     p = '../images/camera_4/' + str(t) + '.png'
#     imwrite(p, img)

topic_rgb = '/camera_1/rgb/image_raw'
window_name = 'image'
i = 0
n = 0
for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
    # if n % 3 == 0:
    img = cv_bridge.imgmsg_to_cv2(msg)
    # img=cv2.ConvertScaleAbs(img,256.0E)
    p = '../images/hpe_stopped_far/camera_1/' + str(i) + '.png'
    # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
    # print(img)
    # print(img.dtype)
    # cv2.imshow(window_name, img)
    # cv2.waitKey(0)
    imwrite(p, img)
    i = i + 1
    # n = n + 1

topic_rgb = '/camera_2/rgb/image_raw'
window_name = 'image'
i = 0
n = 0
for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
    # if n % 3 == 0:
    img = cv_bridge.imgmsg_to_cv2(msg)
    # img=cv2.ConvertScaleAbs(img,256.0E)
    p = '../images/hpe_stopped_far/camera_2/' + str(i) + '.png'
    # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
    # print(img)
    # print(img.dtype)
    # cv2.imshow(window_name, img)
    # cv2.waitKey(0)
    imwrite(p, img)
    i = i + 1
    # n = n + 1

topic_rgb = '/camera_3/rgb/image_raw'
window_name = 'image'
i = 0
n = 0
for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
    # if n % 3 == 0:
    img = cv_bridge.imgmsg_to_cv2(msg)
    # img=cv2.ConvertScaleAbs(img,256.0E)
    p = '../images/hpe_stopped_far/camera_3/' + str(i) + '.png'
    # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
    # print(img)
    # print(img.dtype)
    # cv2.imshow(window_name, img)
    # cv2.waitKey(0)
    imwrite(p, img)
    i = i + 1
    # n = n + 1

topic_rgb = '/camera_4/rgb/image_raw'
window_name = 'image'
i = 0
n = 0
for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
    # if n % 3 == 0:
    img = cv_bridge.imgmsg_to_cv2(msg)
    # img=cv2.ConvertScaleAbs(img,256.0E)
    p = '../images/hpe_stopped_far/camera_4/' + str(i) + '.png'
    # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
    # print(img)
    # print(img.dtype)
    # cv2.imshow(window_name, img)
    # cv2.waitKey(0)
    imwrite(p, img)
    i = i + 1
    # n = n + 1

# DEPTH

# topic_depth = '/camera_1/depth/image_raw'
# window_name = 'image'
# i=0
# for topic_rgb, msg, t in bag.read_messages(topics=[topic_depth]):
#     img = cv_bridge.imgmsg_to_cv2(msg)
#     # cv2.normalize(img, img, 0, 255, cv2.NORM_MINMAX)
#     # np_arr = np.frombuffer(msg.data, np.uint8)
#     # print(np_arr)
#     # img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
#     # cv_image = convertDepthImage32FC1to16UC1(img, scale=10000)  # Better to use tenths of milimeters
#     # cv2.normalize(cv_image, cv_image, 0, 65535, cv2.NORM_MINMAX)
#     p = '../images/with_depth/camera_1_depth/' + str(i) + '.png'
#     cv2.imwrite(p, img)
#     i=i+1