# import time, sys, os
# from ros import rosbag
# import roslib
# from sensor_msgs.msg import CameraInfo, Image, CompressedImage
# from cv_bridge import CvBridge
# from hpe.msg import skeleton, person2D, keypoint2D
# import cv2
# import numpy as np
# import os
# import rospy
#
#
# def create_bag():
#     path = os.path.dirname('/home/daniela/bagfiles/')
#     bagname = path + 'mpi.bag'
#     bag = rosbag.Bag(bagname, 'w')
#
#     # stamp=rospy.Time.now()
#     stamp=time.time()
#     print(stamp)
#
#     bag.close()
#
#     # bag.write('/tf_static', msg_tf_static, stamp, connection_header=tf_static_connection_header)
#
#
# def main():
#     # return
#     rospy.init_node('create_bag')
#
#     create_bag()
#
# if __name__ == '__main__':
#     main()

########################################################################################################################


import time, sys, os

from natsort import natsorted
from ros import rosbag
import roslib
import rospy

roslib.load_manifest('sensor_msgs')
from sensor_msgs.msg import Image, CameraInfo
from utils.load_cam_params import getCamParamsMPI
from hpe.msg import skeleton, person2D, keypoint2D



# from PIL import ImageFile
from PIL import Image as img
import PIL
import cv2
import numpy as np
from itertools import chain


def GetFilesFromDir(dir):
    '''Generates a list of files from the directory'''
    print("Searching directory %s" % dir)
    all = []
    if os.path.exists(dir):
        all_images=natsorted(os.listdir(dir))
        for f in all_images:
            all.append(os.path.join(dir, f))
    return all  # , left_files, right_files


def WriteImagestoBag(bag, all_imgs, camera, cam_params, stamp_list):
    '''Creates a bag file containing stereo image pairs'''

    # try:
    for i in range(len(stamp_list)):
        # Stamp = rospy.rostime.Time.from_sec(time.time())
        Stamp=stamp_list[i]

        #Adding images to rosbag
        print("Adding %s" % all_imgs[i])
        image = img.open(all_imgs[i])
        ros_img = Image()
        ros_img.header.stamp = Stamp
        ros_img.width = image.size[0]
        ros_img.height = image.size[1]
        ros_img.encoding = "rgb8"
        ros_img.header.frame_id = (camera + "_rgb_optical_frame")
        ros_img.data=np.array(image).tobytes()

        bag.write(camera+'/rgb/image_raw', ros_img, Stamp)

        #Adding camera info to rosbag
        camera_info_msg = CameraInfo()
        camera_info_msg.header.stamp=Stamp
        camera_info_msg.width = cam_params[camera]['width']
        camera_info_msg.height = cam_params[camera]['height']
        camera_info_msg.K = list(chain.from_iterable(cam_params[camera]['K'].tolist()))
        camera_info_msg.D = list(cam_params[camera]['dist'])
        camera_info_msg.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]
        camera_info_msg.P = list(chain.from_iterable(np.hstack([np.array(cam_params[camera]['K']), np.zeros((3,1))]).tolist()))
        camera_info_msg.distortion_model = 'plumb_bob'

        bag.write(camera + '/rgb/camera_info', camera_info_msg, Stamp)


def Write2DSkeletontoBag(bag, camera, stamp_list):
    person2d_msg = person2D()
    person2d_msg.header.frame_id = camera + '/rgb/skeleton'

    poses_np = np.load("/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/S1/Seq1/cache/"+ camera +"/2d_pose_mediapipe.npy", allow_pickle=True)
    idx=0
    for pose_np in poses_np.tolist()[0:len(stamp_list)]:
        Stamp=stamp_list[idx]
        person2d_msg.header.stamp = Stamp
        keypoints=pose_np

        keypoints_msgs = []
        for keypoint in keypoints:
            msg = keypoint2D()
            msg.x = keypoint[0]
            msg.y = keypoint[1]
            msg.score = keypoint[2]
            keypoints_msgs.append(msg)
        person2d_msg.keypoints = keypoints_msgs

        bag.write(camera + '/skeleton', person2d_msg, Stamp)
        idx+=1


    return

def CreateBag(args):
    '''Creates the actual bag file by successively adding images'''
    folder="/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/S1/Seq1/imageSequence"
    cameras=['camera_0', 'camera_4', 'camera_5', 'camera_8']
    bag_name="/home/daniela/bagfiles/mpi_slow.bag"

    bag = rosbag.Bag(bag_name, 'w')

    cam_params = getCamParamsMPI('S1', "Seq1", cameras)

    stamp_list=[rospy.rostime.Time.from_sec(time.time())]

    for i in range(500):
        stamp_list.append(stamp_list[-1]+rospy.Duration(0.05))

    for i in range(500):
        stamp_list.append(stamp_list[-1]+rospy.Duration(0.2))

    for camera in cameras:


        all_imgs = GetFilesFromDir(folder+"/"+camera)
        print(len(all_imgs))
        if len(all_imgs) <= 0:
            print("No images found in %s" % args[0])
            exit()

        if len(all_imgs) > 0:
            # CreateMonoBag(all_imgs, args[1])
            Write2DSkeletontoBag(bag, camera, stamp_list)
            WriteImagestoBag(bag,all_imgs, camera, cam_params, stamp_list)


    bag.close()


if __name__ == "__main__":
    CreateBag(sys.argv[1:])
    # if len(sys.argv) == 3:
    #     CreateBag(sys.argv[1:])
    # else:
    #     print("Usage: img2bag imagedir bagfilename")