###########################################################################
# MesobotInference.py
#
# Monterey Bay Aquarium Research Institute   2019
# All rights reserved.
#
# Scaffold for sending and receiving fly_image_t lcm messages in python
# using opencv for display
#
#
# Created:  Oct 2019      Paul Roberts
#
###########################################################################


import time
import sys
from image import image_t
import lcm
import cv2
import numpy as np

LCM_CHANNEL = 'Vim0'
PIXEL_TYPE = 'mono'

"""
FrameTimer - time frame events and estimate and log frame rate.
"""
class FrameTimer:

    def __init__(self, fps=60, log_rate=1, fps_samples=10):
        self.frames = 0
        self.fps_time = 0
        self.log_time = 0
        self.log_rate = log_rate
        self.fps = fps
        self.fps_samples = np.zeros((10,1))
        self.fps_samples_index = 0

    def new_frame(self):
        self.frames += 1
        self.fps_samples = np.roll(self.fps_samples, -1, axis=0)
        self.fps_samples[-1] = time.time()

    def estimate_fps(self):
        fps_est = self.fps_samples[1:-1] - self.fps_samples[0:-2]
        fps_est = np.mean(fps_est)
        if fps_est > 0:
            return 1/fps_est
        else:
            return 0.0

    def fps_start(self):
        self.fps_time = time.time()

    def log_start(self):
        self.log_time = time.time()

    def fps_sleep(self):
        try:
            time.sleep(1/self.fps - (time.time()-self.fps_time))
        except:
            pass

    def log(self):
        if (time.time() - self.log_time) > 1/self.log_rate:
            print('Frames: ' +
                  str(self.frames) + ', FPS: ' + str(self.estimate_fps())
                  )
            self.log_time = time.time()


"""
frame_2_msg - convert a frame from the video (camera) into an lcm message
"""
def frame_2_msg(msg, frame):

    msg.utime = int(time.time() * 1000)
    msg.width = frame.shape[1]
    msg.height = frame.shape[0]
    if PIXEL_TYPE == 'mono':
        msg.data = frame[:, :, 0].tobytes()
    else:
        msg.data = frame.tobytes()

    msg.size = len(msg.data)
    msg.pixelformat = image_t.PIXEL_FORMAT_RGB

    return msg

"""
msg_2_frame - decode an lcm message into a frame
"""
def msg_2_frame(msg):

    frame = np.frombuffer(msg.data, dtype=np.uint8)
    # a few conditions here for image types, many more are possible
    if msg.pixelformat == image_t.PIXEL_FORMAT_RGB:
        bytes_per_pixel = 3
    if msg.pixelformat == image_t.PIXEL_FORMAT_GRAY:
        bytes_per_pixel = 1
    frame = np.reshape(frame, newshape=(msg.height, msg.width, bytes_per_pixel))
    frame = cv2.resize(frame, (960, 540))

    return frame

"""
send_video - Test reading a video file and sending it over lcm
"""
def send_video(video_path, chan=LCM_CHANNEL):

    # setup lcm and image_t
    lc = lcm.LCM()
    img_msg = image_t()

    # open video
    print("Opening: " + video_path)
    cap = cv2.VideoCapture(video_path)

    # Check if video opened successfully
    if (cap.isOpened() == False):
        print("Error opening video stream or file")

    timer = FrameTimer(fps=60)

    # Read until video is completed and report frame rate periodically
    try:
        while cap.isOpened():

            # time the read and send so we know how long to sleep
            timer.fps_start()

            # read and send the frame
            ret, frame = cap.read()
            if ret:
                frame_2_msg(img_msg, frame)
                timer.new_frame()
                lc.publish(chan, img_msg.encode())
            else:
                break

            # sleep to match fps and possibly log frame rate
            timer.fps_sleep()
            timer.log()


    except KeyboardInterrupt:
        pass

def read_log_file(filepath, chan=LCM_CHANNEL):

    log = lcm.EventLog(filepath, "r")

    for event in log:
        print(event.channel)
        if event.channel == chan:
            msg = image_t.decode(event.data)
            print(msg.timestamp)

def read_stereo_images(filepath, left_channel='Vim0', right_channel='Vim1'):

    log = lcm.EventLog(filepath, "r")

    l_msg = None
    r_msg = None

    lc = lcm.LCM()

    timer = FrameTimer(fps=10)
    timer.fps_start()

    try:

        for event in log:
            if event.channel == left_channel:
                l_msg = event.data
            if event.channel == right_channel:
                r_msg = event.data

            # send out messages if have both left and right images
            if (l_msg is not None) and (r_msg is not None):

                lc.publish('DSPL_LEFT', l_msg)
                lc.publish('DSPL_RIGHT', r_msg)
                l_msg = None
                r_msg = None
                timer.new_frame()
                timer.fps_sleep()
                timer.log()
                timer.fps_start()

    except KeyboardInterrupt:
        pass




"""
receive_video - test receiving lcm messages and displaying them with opencv
"""
def receive_video(chan=LCM_CHANNEL):

    def handler(channel, data):
        msg = image_t.decode(data)
        cv2.imshow('Sender', msg_2_frame(msg))
        cv2.waitKey(1)

    timer = FrameTimer()
    lc = lcm.LCM()
    image = image_t()
    subscription = lc.subscribe(chan, handler)

    try:
        while True:
            lc.handle()
            timer.new_frame()
            timer.log()
    except KeyboardInterrupt:
        pass

    lc.unsubscribe(subscription)

"""
Main entry point when called from the command line
"""
if __name__=="__main__":

    print("MesobotInference:")

    if len(sys.argv) < 2:
        print("please input mode [send | receive ] as first argument")
        exit()

    if len(sys.argv) < 3:
        print("Please input LCM channel")
        exit()

    mode = sys.argv[1]
    channel = sys.argv[2]

    if mode == 'receive':
        receive_video(channel)
    else:
        if len(sys.argv) < 4:
            print('Please input video file or lcm log as third argument for send mode')
            exit()

        if mode == 'send':
            send_video(sys.argv[3], channel)
            exit()

        if mode == 'load_stereo':
            read_stereo_images(sys.argv[3])