import rospy, rosbag
import numpy as np
import cv2
import pickle
import argparse


def image_to_numpy(msg):
    """
    Convert a sensor_msgs/CompressedImage to numpy array
    """
    img_np_arr = np.fromstring(msg.data, np.uint8)
    img = cv2.imdecode(img_np_arr, cv2.IMREAD_COLOR)
    img = np.roll(img, 2, axis=-1) # To get the colors correct

    return img


def main(args):
    """
    Bag to state-action pairs, stored in a pickle file
    """
    bag = rosbag.Bag(args.bagfile)

    # Initialize your current action to be 0s
    cur_action = (0., 0.)

    actions = []
    images = []
    for topic, msg, t in bag.read_messages(topics=[args.img_topic, args.action_topic]):
        if topic == args.action_topic:
            # Change this line if your action msg is not of type duckietown_msgs/WheelsCmdStamped
            cur_action = (msg.vel_left, msg.vel_right)
        if topic == args.img_topic:
            actions.append(cur_action)

            img = image_to_numpy(msg)
            images.append(img)

    actions = np.array(actions)
    images = np.array(images)

    with open(args.outfile, "wb") as f:
        pickle.dump((images, actions), f)

def load(args):
    """
    Test to make sure you have the right amount of sa pairs
    """
    with open(args.outfile, "rb") as f:
        img, act = pickle.load(f)

    # Set show_proof to True if you'd like to see the image
    if args.show_proof:
        from PIL import Image
        im = Image.fromarray(img[-1])
        im.save('test.png')

        print('Last 100 Actions', act[-100:])

    print(img.shape, act.shape)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Rosbag->Numpy Synchronizer')
    parser.add_argument('--bagfile', type=str, help='bagfile name')
    parser.add_argument('--outfile', type=str, help='outfile name')
    parser.add_argument('--img-topic', type=str, help='image topic')
    parser.add_argument('--action-topic', type=str, help='action topic')
    parser.add_argument('--show-proof', help='Show proof that it worked', action='store_true')

    args = parser.parse_args()

    main(args)
    load(args)