###########################################################################
# tools.py
#
# Monterey Bay Aquarium Research Institute   2019
# All rights reserved.
#
# General tools for examples
#
#
# Created:  Dec 2019      Paul Roberts
#
###########################################################################

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

"""
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, msg.utime