###########################################################################
# MLHandler.py
#
# Modified Version of lcm_image_test.py from Ben and Jon at CVision
#
# The modifications support:
#   (1) variable input image size
#   (2) Variable image display size
#   (3) Variable image format (mono or rgb)
#   (4) Real-time display of images using either qt signals or opencv
#
# Adapted:  Nov 2019      Paul Roberts
#
###########################################################################

import datetime
import cv2
import os
import numpy as np
import time
import logging
from collections import deque
from image import image_t
from mbari import ext_target_t, ext_target_wheader_t
from cvision import bounding_box_t
from pyqtgraph.Qt import QtCore
from stereo_calc import StereoCalc
from lcm import LCM
from annotation import Annotation

logging.basicConfig(
    format='%(asctime)s %(levelname)-8s %(message)s',
    level=logging.DEBUG,
    datefmt='%Y-%m-%d %H:%M:%S',
)
logger = logging.getLogger(__name__)


class Handler(QtCore.QThread):

    stereoFrame = QtCore.Signal(object)

    @classmethod
    def init_class(cls, args):
        fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
        # cls.writer = cv2.VideoWriter(video_path, fourcc, 10, (1032, 1544))
        cls.target_class_name = args.target_class
        cls.display_width = args.display_width
        cls.display_height = args.display_height
        cls.box_width = 120
        cls.box_height = 120
        cls.imgs = [deque([], maxlen=10), deque([], maxlen=10)]
        cls.boxes = []
        cls.num_processed = 0
        cls.num_boxes = 0
        cls.num_targets = 0
        cls.score_threshold = 0.25
        cls.box_match_threshold = args.default_box_dev
        cls.match_classes = True
        cls.match_method = None
        cls.save_ml_rois = False
        cls.save_kcf_rois = False
        cls.kcf_coords = None
        cls.show_bg_mask = False
        cls.show_ml = False
        cls.kcf_okay = False
        cls.stereo_calc = StereoCalc()
        cls.match_methods = []
        cls.img_buffer = np.zeros((cls.display_height * 2, cls.display_width, 3), dtype='uint8')
        cls.box_buffer = [[], []]
        cls.target_buffer = [[], []]
        cls.start_time = datetime.datetime.now()

        # setup KCF tracker
        cls.kcf_trackers = []
        cls.kcf_bboxes = []

        # Setup SimpleBlobDetector parameters.
        cls.params = cv2.SimpleBlobDetector_Params()

        cls.bg_detector_slow = cv2.createBackgroundSubtractorKNN(500)
        cls.bg_detector_mid = cv2.createBackgroundSubtractorKNN(50)
        cls.bg_detector_fast = cv2.createBackgroundSubtractorKNN(5)

        # Change thresholds
        cls.params.minThreshold = 5
        cls.params.maxThreshold = 100

        # Filter by Area.
        cls.params.filterByArea = True
        cls.params.minArea = 50

        # Filter by Circularity
        cls.params.filterByCircularity = False
        cls.params.minCircularity = 0

        # Filter by Convexity
        cls.params.filterByConvexity = False
        cls.params.minConvexity = 0

        # Filter by Inertia
        cls.params.filterByInertia = False
        cls.params.minInertiaRatio = 0

        cls.last_save_utime = int(time.time())

        # Create a detector with the parameters
        cls.detector = cv2.SimpleBlobDetector_create(cls.params)


    @classmethod
    def set_lcm_dirs(cls, folder, file):
        cls.lcm_log_file = file
        cls.lcm_log_folder = folder

    @classmethod
    def set_data_dirs(cls, left_dir, right_dir):

        # Create dirs for storing output
        cls.left_data_dir = left_dir + str(int(time.time()))
        cls.right_data_dir = right_dir + str(int(time.time()))
        if not os.path.exists(Handler.left_data_dir):
            os.makedirs(os.path.join(Handler.left_data_dir, 'xml'))
            os.makedirs(os.path.join(Handler.left_data_dir, 'images'))
            os.makedirs(os.path.join(Handler.left_data_dir, 'boxed_images'))
            os.makedirs(os.path.join(Handler.left_data_dir, 'rois'))
        if not os.path.exists(Handler.right_data_dir):
            os.makedirs(os.path.join(Handler.right_data_dir, 'xml'))
            os.makedirs(os.path.join(Handler.right_data_dir, 'images'))
            os.makedirs(os.path.join(Handler.right_data_dir, 'boxed_images'))
            os.makedirs(os.path.join(Handler.right_data_dir, 'rois'))

    @classmethod
    def set_box_width(cls, width):
        cls.box_width = width

    @classmethod
    def set_box_height(cls, height):
        cls.box_height = height

    @classmethod
    def set_score_threshold(cls, score_threshold):
        logger.info(score_threshold)
        cls.score_threshold = float(score_threshold)/100

    @classmethod
    def set_match_classes(cls, should_match):
        logger.info(should_match)
        cls.match_classes = should_match

    @classmethod
    def set_show_bg_mask(cls, show_bg_mask):
        logger.info(show_bg_mask)
        cls.show_bg_mask = show_bg_mask

    @classmethod
    def set_show_ml(cls, show_ml):
        logger.info(show_ml)
        cls.show_ml = show_ml

    @classmethod
    def set_save_kcf(cls, save_kcf):
        logger.info(save_kcf)
        cls.save_kcf_rois = save_kcf

    @classmethod
    def set_save_ml(cls, save_ml):
        logger.info(save_ml)
        cls.save_ml_rois = save_ml

    @classmethod
    def set_match_method(cls, method):
        cls.match_method = method
        logger.info(method)

    @classmethod
    def set_target_class(cls, target_class):
        cls.target_class_name = target_class
        logger.info(cls.target_class_name)

    @classmethod
    def set_box_deviation(cls, dev):
        logger.info(dev)
        cls.box_match_threshold = dev

    def __init__(self, index):
        QtCore.QThread.__init__(self)
        self.index = index
        self.num_images = 0
        self.init_kcf = False


    def initKCF(self, coords):
        # coords are for hstack, map to vstack
        x = coords[0]
        y = coords[1]
        if x > self.display_width:
            x = x - self.display_width
            y = self.display_height*2 - y
        else:
            y = self.display_height - y
        bbox = (x - self.box_width/2, y - self.box_height/2, self.box_width, self.box_height)

        # check that coords are not on existing box
        # if they are, delete it
        delete_box = -1
        for ind, box in enumerate(Handler.kcf_bboxes):
            if box[0] < x and box[0] + box[2] > x:
                if box[1] < y and box[1] + box[3] > y:
                    delete_box = ind
        if delete_box >= 0:
            Handler.kcf_trackers.pop(delete_box)
            Handler.kcf_bboxes.pop(delete_box)
        else:
            trkr = cv2.TrackerKCF_create()
            trkr.init(Handler.img_buffer, bbox)
            Handler.kcf_trackers.append(trkr)
            Handler.kcf_bboxes.append(bbox)

    def unpack_image(self, msg):

        pix = np.frombuffer(msg.data, dtype='uint8')

        if msg.pixelformat != image_t.PIXEL_FORMAT_GRAY:
            # color image, unpack and resize
            img = pix.reshape(msg.height, msg.width, 3)
        else:
            img = np.tile(pix.reshape(msg.height, msg.width, 1), (1, 1, 3))

        img = cv2.resize(img, (self.display_width, self.display_height))

        return img

    def get_matching_classes(self, left_boxes, right_boxes):

        class_matches = [[]]
        for i in range(len(left_boxes)):
            for j in range(len(right_boxes)):
                if left_boxes[i][3] == right_boxes[j][3]:
                    class_matches.append([left_boxes[i], right_boxes[j]])

        return class_matches

    def get_all_box_pairs(self, left_boxes, right_boxes):

        box_pairs = [[]]
        for i in range(len(left_boxes)):
            for j in range(len(right_boxes)):
                box_pairs.append([left_boxes[i], right_boxes[j]])

        return box_pairs

    def box_match_merge(self, box_pairs, a_thresh=0.10):

        for match in box_pairs:
            if len(match) > 0:

                l_w = abs(match[0][1][0] - match[0][0][0])
                l_h = abs(match[0][1][1] - match[0][1][0])
                r_w = abs(match[1][1][0] - match[1][0][0])
                r_h = abs(match[1][1][1] - match[1][1][0])

                if abs(l_w - r_w) < a_thresh * l_w and abs(l_h - r_h) < a_thresh * l_h:
                    return match

        return None

        if len(left_boxes) != 0:
            print('Target Found: ' + self.target_class_name + str([len(left_boxes), len(right_boxes)]))

    def epipolar_error_match(self, box_pairs):
        #TODO: call stereo calc code with points from boxes
        pass

    def merge_target_boxes(self, left_boxes, right_boxes, a_thresh=0.10):

        if Handler.match_classes:
            box_pairs = self.get_matching_classes(left_boxes, right_boxes)
        else:
            box_pairs = self.get_all_box_pairs(left_boxes, right_boxes)

        if Handler.match_method == 'Box Similarity':
            return self.box_match_merge(box_pairs, Handler.box_match_threshold)
        if Handler.match_method == 'Epipolar Error':
            return self.epipolar_error_match(box_pairs)

    def scale_points(self, p0, p1, scale0, scale1):
        scale = [scale0, scale1]
        p0_out = [0, 0]
        p1_out = [0, 0]
        for i in range(0, len(p0)):
            p0_out[i] = int(p0[i] * scale[i])
            p1_out[i] = int(p1[i] * scale[i])

        return tuple(p0_out), tuple(p1_out)

    def findROIS(self, image, low=100, high=200):
        """
        edges = cv2.Canny(image, low, high)
        kernel = np.ones((3, 3), np.uint8)
        edges_filtered = cv2.morphologyEx(edges, cv2.MORPH_OPEN, kernel)
        contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # find the contour with the largest area
        max_area_index = 0
        max_area = -1
        cont_image = image.copy()
        for ind, cont in enumerate(contours):
            if cv2.contourArea(cont) > 50:
                cv2.drawContours(cont_image, [cont], 0, (0, 0, 255), 1, cv2.LINE_8, hierarchy, 0)

        return cont_image
        """
        #image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        #keypoints = self.detector.detect(image)

        # Draw detected blobs as red circles.
        # cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS ensures the size of the circle corresponds to the size of blob
        #im_with_keypoints = cv2.drawKeypoints(image, keypoints, np.array([]), (0, 0, 255),
        #                                      cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

        fgMask = self.bg_detector_slow.apply(image)

        contours, hierarchy = cv2.findContours(fgMask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # find the contour with the largest area
        max_area_index = 0
        max_area = -1
        cont_image = image.copy()
        for ind, cont in enumerate(contours):
            if cv2.contourArea(cont) > 150:
                cv2.drawContours(cont_image, [cont], 0, (0, 0, 255), 1, cv2.LINE_8, hierarchy, 0)

        return cont_image

    def targtrasn(self, msg):

        msg_out = ext_target_wheader_t()
        msg_out.conf_left = msg.conf_left
        msg_out.conf_right = msg.conf_right
        msg_out.left_class_name = msg.left_class_name
        msg_out.right_class_name = msg.right_class_name
        msg_out.left_utime = msg.left_utime
        msg_out.right_utime = msg.right_utime
        msg_out.left_x = msg.left_x
        msg_out.right_x = msg.right_x
        msg_out.left_y = msg.left_y
        msg_out.right_y = msg.right_y

    def pub_ext_target_3(self, left_x, left_y, right_x, right_y, left_utime, right_utime, right_class_name, left_class_name, conf_left, conf_right):
        msg_out = ext_target_wheader_t()
        msg_out.conf_left = 100*conf_left
        msg_out.conf_right = 100*conf_right
        msg_out.left_class_name = left_class_name
        msg_out.right_class_name = right_class_name
        msg_out.left_utime = left_utime
        msg_out.right_utime = right_utime
        msg_out.left_x = left_x
        msg_out.right_x = right_x
        msg_out.left_y = left_y
        msg_out.right_y = right_y

        lc = LCM()
        lc.publish('EXT_TARGET_3', msg_out.encode())

    def __call__(self, channel, data):
        if channel == 'BOUNDING_BOX':
            if Handler.num_boxes % 100 == 0:
                logger.info(f"Received {self.num_boxes} boxes...")
            Handler.num_boxes += 1
            box = bounding_box_t.decode(data)
            pt0 = (round(box.left), round(box.top))
            pt1 = (round(box.left + box.width), round(box.top + box.height))
            Handler.box_buffer[box.channel].append((pt0, pt1, box.utime, box.class_name, box.scores))
        elif channel == 'EXT_TARGET_2':
            if Handler.num_targets % 100 == 0:
                logger.info(f"Received {self.num_targets} targets...")
            Handler.num_targets += 1
            target = ext_target_t.decode(data)
            left = (round(target.left_x), round(target.left_y))
            right = (round(target.right_x), round(target.right_y))
            Handler.target_buffer[0].append((left, target.left_utime))
            Handler.target_buffer[1].append((right, target.right_utime))

            if Handler.match_method == 'Ml Output':

                lc = LCM()
                lc.publish('EXT_TARGET_3', self.targtrasn(target).encode())
        else:
            if self.num_images % 100 == 0:
                logger.info(f"Received {self.num_images} images on {self.index}...")
            self.num_images += 1
            Handler.imgs[self.index].append(image_t.decode(data))
            if (len(Handler.imgs[0]) > 5) and (len(Handler.imgs[1]) > 5):
                left = Handler.imgs[0].popleft()
                right = Handler.imgs[1].popleft()
                while True:
                    delta = abs(left.utime - right.utime)
                    if delta > 1000:
                        if left.utime < right.utime:
                            logger.info("Dropping a left frame for sync...")
                            left = Handler.imgs[0].popleft()
                        elif right.utime < left.utime:
                            right = Handler.imgs[1].popleft()
                            logger.info("Dropping a right frame for sync...")
                    else:
                        break

                # assume stereo images have the same size
                scale_x = self.display_width / left.width
                scale_y = self.display_height / left.height

                Handler.img_buffer[:self.display_height, :, :] = self.unpack_image(left)
                Handler.img_buffer[self.display_height:, :, :] = self.unpack_image(right)

                left_boxes = []
                right_boxes = []


                # Get copies of images before drawing on then for annotations
                if self.save_ml_rois:
                    left_ann = Annotation(image=Handler.img_buffer[:self.display_height, :, :].copy(),
                                            filename='ImageHandler-Left-' + str(left.utime),
                                            folder=Handler.left_data_dir,
                                            utime=left.utime,
                                            camera_name='left'
                                            )

                    right_ann = Annotation(image=Handler.img_buffer[self.display_height:, :, :].copy(),
                                            filename='ImageHandler-Right-' + str(right.utime),
                                            folder=Handler.right_data_dir,
                                            utime=right.utime,
                                            camera_name='right'
                                            )

                # run Canny on image buffer and draw contours
                if self.show_bg_mask:
                    Handler.img_buffer = self.findROIS(Handler.img_buffer)

                for box in Handler.box_buffer[0]:
                    pt0, pt1, utime, class_name, scores = box
                    if np.max(scores) > Handler.score_threshold:
                        pt0, pt1 = self.scale_points(pt0, pt1, scale_x, scale_y)
                        if class_name == self.target_class_name:
                            left_boxes.append(box)
                        if utime == left.utime:
                            if self.save_ml_rois:
                                left_ann.add_box(box)
                            if self.show_ml:
                                cv2.rectangle(
                                    Handler.img_buffer[:self.display_height, :, :], tuple(pt0), tuple(pt1), (0, 255, 0), 2
                                )
                                cv2.putText(Handler.img_buffer[:self.display_height, :, :], class_name, (pt0[0], pt1[1]),
                                            cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 255))

                for box in Handler.box_buffer[1]:
                    pt0, pt1, utime, class_name, scores = box
                    if np.max(scores) > Handler.score_threshold:
                        pt0, pt1 = self.scale_points(pt0, pt1, scale_x, scale_y)
                        if class_name == self.target_class_name:
                            right_boxes.append(box)
                        if utime == right.utime:
                            if self.save_ml_rois:
                                right_ann.add_box(box)
                            if self.show_ml:
                                cv2.rectangle(
                                    Handler.img_buffer[self.display_height:, :, :], tuple(pt0), tuple(pt1), (0, 255, 0), 2
                                )
                                cv2.putText(Handler.img_buffer[self.display_height:, :, :], class_name, (pt0[0], pt1[1]),
                                            cv2.FONT_HERSHEY_COMPLEX, 0.5,
                                            (0, 255, 255))
                for target in Handler.target_buffer[0]:
                    pt, utime = target
                    pt, pt = self.scale_points(pt, pt, scale_x, scale_y)
                    if utime == left.utime:
                        cv2.circle(Handler.img_buffer[:self.display_height, :, :], tuple(pt), 3, (255, 0, 0), 5)
                for target in Handler.target_buffer[1]:
                    pt, utime = target
                    pt, pt = self.scale_points(pt, pt, scale_x, scale_y)
                    if utime == right.utime:
                        cv2.circle(Handler.img_buffer[self.display_height:, :, :], tuple(pt), 3, (255, 0, 0), 5)

                fused = self.merge_target_boxes(left_boxes, right_boxes)


                # save out rois and xmls once a second
                if self.save_ml_rois and (int(time.time()) - self.last_save_utime > 0):
                    if len(left_ann.bboxes) > 0:
                        left_ann.save_xml()
                        left_ann.save_image()
                        left_ann.save_rois()
                    if len(right_ann.bboxes) > 0:
                        right_ann.save_xml()
                        right_ann.save_image()
                        right_ann.save_rois()

                    self.last_save_utime = int(time.time())

                for ind, tracker in enumerate(Handler.kcf_trackers):
                    ok, bbox = tracker.update(Handler.img_buffer)
                    Handler.kcf_bboxes[ind] = bbox
                    cv2.rectangle(
                        Handler.img_buffer, tuple([int(bbox[0]), int(bbox[1])]), tuple([int(bbox[0] + bbox[2]), int(bbox[1]+bbox[3])]), (200, 255, 0), 2
                    )

                    if self.save_kcf_rois:
                        print('save kcf here')



                if fused is not None:
                    pt00, pt01 = self.scale_points(fused[0][0], fused[0][1], scale_x, scale_y)
                    pt10, pt11 = self.scale_points(fused[1][0], fused[1][1], scale_x, scale_y)
                    # Compute  the centroid
                    left_x = pt00[0] + (pt01[0] - pt00[0]) / 2
                    left_y = pt00[1] + (pt01[1] - pt00[1]) / 2
                    right_x = pt10[0] + (pt11[0] - pt10[0]) / 2
                    right_y = pt10[1] + (pt11[1] - pt10[1]) / 2
                    conf_left = max(fused[0][-1])
                    conf_right = max(fused[1][-1])
                    left_class_name = fused[0][-2]
                    right_class_name = fused[1][-2]
                    left_utime = fused[0][2]
                    right_utime = fused[1][2]

                    self.pub_ext_target_3(left_x, left_y, right_x, right_y, left_utime, right_utime, right_class_name,
                                     left_class_name, conf_left, conf_right)

                    cv2.rectangle(
                        Handler.img_buffer[:self.display_height, :, :], pt00, pt01, (0, 255, 255), 2
                    )
                    cv2.rectangle(
                        Handler.img_buffer[self.display_height:, :, :], pt10, pt11, (0, 255, 255), 2
                    )

                # Handler.writer.write(Handler.img_buffer)
                #cv2.imshow('ML-Output', Handler.img_buffer)
                #cv2.waitKey(1)
                self.stereoFrame.emit(Handler.img_buffer)
                Handler.num_processed += 1
                if Handler.num_processed % 100 == 0:
                    elapsed = datetime.datetime.now() - Handler.start_time
                    elapsed = elapsed.total_seconds()
                    logger.info(
                        f"Processed {Handler.num_processed} stereo "
                        f"pairs in {elapsed:.2f} seconds "
                        f"({Handler.num_processed / elapsed:.2f} Hz)..."
                    )
                Handler.box_buffer = [
                    [b for b in Handler.box_buffer[0] if b[2] > left.utime],
                    [b for b in Handler.box_buffer[1] if b[2] > right.utime],
                ]
                Handler.target_buffer = [
                    [t for t in Handler.target_buffer[0] if t[1] > left.utime],
                    [t for t in Handler.target_buffer[1] if t[1] > right.utime],
                ]