import base64
from datetime import datetime
from collections import deque
import cv2
import numpy as np
from edge_engine.ai.model.modelwraper import ModelWrapper
from edge_engine.common.logsetup import logger
from expiringdict import ExpiringDict
from scripts.common.constants import JanusDeploymentConstants
from scripts.utils.centroidtracker import CentroidTracker
from scripts.utils.helpers import box_iou2
from scripts.utils.image_utils import draw_circles_on_frame
from scripts.utils.tracker import Tracker
from scripts.utils.yolov5_trt import YoloV5TRT
from scipy.optimize import linear_sum_assignment as linear_assignment
from edge_engine.common.config import EDGE_CONFIG, LAST_COUNT, MONGO_URI_ADD_EVENTLOGS
import requests
from uuid import uuid1

class CementBagCounter(ModelWrapper):

    def __init__(self, config, model_config, pubs, device_id):
        super().__init__()
        """
        init function
        """
        logger.info("Initializing Cement Bag Counter")
        self.config = config["config"]
        logger.info(f"Config Received {self.config}")
        self.device_id = device_id
        self.rtp = pubs.rtp_write
        self.frame_skip = self.config.get('frame_skip', False)
        engine_file_path = "data/jk_v5_cam_44.engine"
        self.classes = {0: 'cement_bag'}
        self.yolo_v5_wrapper = YoloV5TRT(engine_file_path, model_config.get('conf_thresh', 0.7),
                                         model_config.get('iou_thresh', 0.15))

        # TRT Additions stop

        self.print_eu_dist = model_config.get('print_eu_dist', 200)
        self.ct1 = CentroidTracker(maxDisappeared=5)
        self.ct2 = CentroidTracker(maxDisappeared=5)

        self.count = int(LAST_COUNT)
        self.tracker_list = []
        self.max_age = 3
        self.min_hits = 0
        self.track_id_list = deque([str(i) for i in range(1, 50)])
        self.prev_annotation = []
        self.initial_object_position = None
        self.uncounted_objects = ExpiringDict(max_len=model_config.get("uncounted_obj_length", 50),
                                              max_age_seconds=model_config.get("uncounted_obj_age", 60))
        self.janus_metadata = ExpiringDict(max_age_seconds=120, max_len=1)
        # self.camera_details = EDGE_CONFIG[self.device_id]
        self.frame_count = 0
        self.frame = None
        self.centroid_distance = 10


    def _pre_process(self, x):
        """
        Do preprocessing here, if any
        :param x: payload
        :return: payload
        """
        return x

    def _post_process(self, x):
        """
        Apply post processing here, if any
        :param x: payload
        :return: payload
        """
        self.rtp.publish(x)  # video stream
        logger.info("Pushed frame")
        return x

    def kalman_tracker(
            self,
            bboxs,
            img,
    ):

        z_box = bboxs
        x_box = []
        logger.debug("Using Kalman Tracker to Track Objects")
        if len(self.tracker_list) > 0:
            logger.debug(f"length of tracker list is greater than zero")
            for trk in self.tracker_list:
                logger.debug("Appending tracker list")
                x_box.append(trk.box)
        logger.debug("Assigning detections to trackers")
        matched, unmatched_dets, unmatched_trks = self.assign_detections_to_trackers(x_box, z_box, iou_thrd=0.15)

        # Deal with matched detections
        if matched.size > 0:
            logger.debug("Deal with matched detections")
            for trk_idx, det_idx in matched:
                z = z_box[det_idx]
                z = np.expand_dims(z, axis=0).T
                tmp_trk = self.tracker_list[trk_idx]
                tmp_trk.kalman_filter(z)
                xx = tmp_trk.x_state.T[0].tolist()
                xx = [xx[0], xx[2], xx[4], xx[6]]
                x_box[trk_idx] = xx
                tmp_trk.box = xx
                tmp_trk.hits += 1

        # Deal with unmatched detections
        if len(unmatched_dets) > 0:
            logger.debug("Deal with unmatched detections")
            for idx in unmatched_dets:
                z = z_box[idx]
                z = np.expand_dims(z, axis=0).T
                tmp_trk = Tracker()  # Create a new tracker
                x = np.array([[z[0], 0, z[1], 0, z[2], 0, z[3], 0]]).T
                tmp_trk.x_state = x
                tmp_trk.predict_only()
                xx = tmp_trk.x_state
                xx = xx.T[0].tolist()
                xx = [xx[0], xx[2], xx[4], xx[6]]
                tmp_trk.box = xx
                tmp_trk.id = self.track_id_list.popleft()  # assign an ID for the tracker

                self.tracker_list.append(tmp_trk)
                x_box.append(xx)

        # Deal with unmatched tracks
        if len(unmatched_trks) > 0:
            logger.debug("Deal with unmatched tracks")
            for trk_idx in unmatched_trks:
                tmp_trk = self.tracker_list[trk_idx]
                tmp_trk.no_losses += 1
                tmp_trk.predict_only()
                xx = tmp_trk.x_state
                xx = xx.T[0].tolist()
                xx = [xx[0], xx[2], xx[4], xx[6]]
                tmp_trk.box = xx
                x_box[trk_idx] = xx

        # The list of tracks to be annotated
        good_tracker_list = []
        objects = []
        boxs = []
        for trk in self.tracker_list:
            logger.debug("Getting the list of trackers to be annotated")
            if (trk.hits >= self.min_hits) and (trk.no_losses <= self.max_age):
                good_tracker_list.append(trk)
                x_cv2 = trk.box
                left, top, right, bottom = x_cv2[1], x_cv2[0], x_cv2[3], x_cv2[2]
                centroid = [int(left + ((right - left) / 2)), bottom]
                objects.append([int(trk.id), centroid])
                boxs.append(x_cv2)

        deleted_tracks = filter(lambda _x: _x.no_losses > self.max_age, self.tracker_list)

        for trk in deleted_tracks:
            self.track_id_list.append(trk.id)

        self.tracker_list = [x for x in self.tracker_list if x.no_losses <= self.max_age]

        return img, objects, boxs

    @staticmethod
    def assign_detections_to_trackers(
            trackers,
            detections,
            iou_thrd=0.15,
    ):
        """
        From current list of trackers and new detections, output matched detections,
        un matched trackers, unmatched detections.
        """
        iou_mat = np.zeros((len(trackers), len(detections)), dtype=np.float32)
        for t, trk in enumerate(trackers):
            logger.debug("Enumerating through trackers")
            for d, det in enumerate(detections):
                logger.debug("Enumerating through detections")
                iou_mat[t, d] = box_iou2(trk, det)
        logger.debug("Linear assignment")
        matched_idx = linear_assignment(-iou_mat)
        row, col = matched_idx
        matched_idx = np.concatenate((row.reshape(-1, 1), col.reshape(-1, 1)), axis=1)
        logger.debug("Linear assignment done")
        unmatched_trackers, unmatched_detections = [], []
        for t, trk in enumerate(trackers):
            if t not in matched_idx[:, 0]:
                unmatched_trackers.append(t)

        for d, det in enumerate(detections):
            if d not in matched_idx[:, 1]:
                unmatched_detections.append(d)

        matches = []

        for m in matched_idx:
            if iou_mat[m[0], m[1]] < iou_thrd:
                unmatched_trackers.append(m[0])
                unmatched_detections.append(m[1])
            else:
                matches.append(m.reshape(1, 2))

        if len(matches) == 0:
            matches = np.empty((0, 2), dtype=int)
        else:
            matches = np.concatenate(matches, axis=0)
        logger.debug("leaving assign detections to trackers function")
        return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

    def get_line_coordinates(self):
        """
        Get the line coordinates from the deployment JSON
        """
        if not self.janus_metadata.get('metadata'):
            self.janus_metadata['metadata'] = EDGE_CONFIG["extra_fields"]
        _coordinates = []
        for data in self.janus_metadata["metadata"]:
            for elem in data:
                if data[elem] in JanusDeploymentConstants.LINE_COORDINATES:
                    _coordinates.append(data["value"])
                if data[elem] == JanusDeploymentConstants.ALIGNMENT_KEY:
                    _alignment = data["value"]

        return _alignment, _coordinates

    def line_point_position(self, point):
        """
        Get the position of point w.r.t. the line
        :param point: point to be compared
        :return: boolean
        """
        _alignment, line_coordinates = self.get_line_coordinates()

        assert len(line_coordinates) == 4, "Line coordinates variable is invalid"
        assert len(point) == 2, "Point variable is invalid"

        _slope = (line_coordinates[3] - line_coordinates[1]) / (line_coordinates[2] - line_coordinates[0])
        _point_equation_value = point[1] - line_coordinates[1] - _slope * (point[0] - line_coordinates[0])
        if _point_equation_value > 0:
            return True
        else:
            return False


    def validate_point_position(self, point):
        """
        Validate the position of the point w.r.t. the line
        :param point: centroid
        :return: bool
        """
        _alignment, line_coordinates = self.get_line_coordinates()
        assert _alignment in [JanusDeploymentConstants.VERTICAL, JanusDeploymentConstants.HORIZONTAL], \
            "Invalid alignment variable"
        if _alignment == JanusDeploymentConstants.VERTICAL:
            line_y2 = line_coordinates[3]
            line_y1 = line_coordinates[1]
            if line_y1 < point[1] < line_y2 or line_y2 < point[1] < line_y1:
                return True
            else:
                return False
        else:
            line_x2 = line_coordinates[2]
            line_x1 = line_coordinates[0]
            if line_x1 < point[0] < line_x2 or line_x2 < point[0] < line_x1:

                return True
            else:
                return False

    def update_bag_count(self, frame, detection_objects, class_name, detections):
        """
        Maintains the bag counts
        :param frame: image
        :param detection_objects: detection object having object id and centroids
        """
        logger.debug(f"Updating Cement Bag Count and detection objects are {detection_objects}")
        for (object_id, det, class_detected) in zip(
                detection_objects, detections, class_name
        ):
            centroid = object_id[1]
            object_id = object_id[0]
            frame = draw_circles_on_frame(
                frame, centroid, radius=10, color=(0, 0, 255), thickness=-1
            )
            if self.validate_point_position(centroid):

                if not isinstance(self.initial_object_position, bool):
                    logger.debug("Initializing the initial object position")
                    # self.initial_object_position = self.line_point_position(point=centroid)
                    self.initial_object_position = True

                _point_position = self.line_point_position(point=centroid)
                # Check point in the same side as the initial object
                if _point_position == self.initial_object_position:
                    # Check the object is not already counted
                    if object_id not in self.uncounted_objects:
                        self.uncounted_objects[object_id] = centroid

                elif object_id in self.uncounted_objects:
                    self.uncounted_objects.pop(object_id, None)
                    self.count += 1
                    img2 = frame.copy()
                    img2 = cv2.resize(img2, (64, 64))
                    img2 = cv2.imencode('.jpg', img2)[1]
                    bs64 = str(base64.b64encode(img2).decode('utf-8'))
                    bs64 += "data:image/jpeg;base64,"
                    json = {"eventId": str(uuid1()).split('-')[0], "cameraId": EDGE_CONFIG['deviceId'], "cameraName": EDGE_CONFIG['deviceId'], "timestamp": datetime.isoformat(datetime.now()), "frame": bs64, "cement_bag_count": str(self.count)}
                    try:
                        requests.post(url=MONGO_URI_ADD_EVENTLOGS, json=json)
                    except Exception as e:
                        logger.info(f"Failed to push count to server: {e}")
                    logger.info(f"Count: {self.count}")
                    frame = draw_circles_on_frame(
                        frame, centroid, radius=10, color=(0, 255, 0), thickness=-1
                    )

                else:
                    frame = draw_circles_on_frame(
                        frame, centroid, radius=10, color=(0, 255, 0), thickness=-1
                    )

        return frame

    def draw_line_over_image(self, frame, color=(255, 255, 255)):
        """
        Draws line over the counting line
        :param frame: frame for
        :param color:
        :return:
        """
        _alignment, line_coordinates = self.get_line_coordinates()
        assert len(line_coordinates) == 4, "Line coordinates variable is invalid"

        self.drawline(
            frame,
            (line_coordinates[0], line_coordinates[1]),
            (line_coordinates[2], line_coordinates[3]),
            color,
            thickness=3,
        )
        return frame

    @staticmethod
    def drawline(img, pt1, pt2, color, thickness=1, style="dotted", gap=20):
        dist = ((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2) ** 0.5
        pts = []
        for i in np.arange(0, dist, gap):
            r = i / dist
            x = int((pt1[0] * (1 - r) + pt2[0] * r) + 0.5)
            y = int((pt1[1] * (1 - r) + pt2[1] * r) + 0.5)
            p = (x, y)
            pts.append(p)

        if style == "dotted":
            for p in pts:
                cv2.circle(img, p, thickness, color, -1)
        else:
            s = pts[0]
            e = pts[0]
            i = 0
            for p in pts:
                s = e
                e = p
                if i % 2 == 1:
                    cv2.line(img, s, e, color, thickness)
                i += 1

    def overlapping_removal(self, frame, dets):
        h, w, _ = frame.shape
        removable_detections = []

        centroid_distance_in_pixels = (self.centroid_distance * w) / 100
        logger.debug(f"centroid_distance_in_pixels : {centroid_distance_in_pixels}")
        if dets:
            for j in dets:
                centroid_1 = j["points"][0] + (j["points"][2] - j["points"][0]) / 2
                conf_1 = j["conf"]
                for k in dets:
                    if j != k:

                        centroid_2 = k["points"][0] + (k["points"][2] - k["points"][0]) / 2
                        conf_2 = k["conf"]
                        centroid_distance = abs(centroid_2 - centroid_1)
                        logger.debug(f"centroid distance : {centroid_distance}")
                        if centroid_distance < centroid_distance_in_pixels:
                            if conf_1 > conf_2:
                                if k not in removable_detections:
                                    removable_detections.append(k)
                            else:
                                if j not in removable_detections:
                                    removable_detections.append(j)
            for elem in removable_detections:
                frame = cv2.rectangle(frame, (elem["points"][0], elem["points"][1]),
                                      (elem["points"][2], elem["points"][3]),
                                      (255, 255, 0), 2)

                dets.remove(elem)

            for x in dets:
                frame = cv2.rectangle(frame, (x["points"][0], x["points"][1]), (x["points"][2], x["points"][3]),
                                      (255, 255, 255), 2)
        logger.debug(f"Dets after overlapping removal: {dets}")
        return dets, frame

    def inference(
            self,
            frame):
        # TRT Additions start
        result_boxes, result_scores, result_classid = self.yolo_v5_wrapper.infer(frame)
        dets = [{"points": list(points), "conf": conf, "class": self.classes.get(class_id)} for points, conf, class_id
                in
                zip(result_boxes, result_scores, result_classid)]
        # TRT Additions stop


        class_name = list()
        bboxs = []

        if dets:
            # adding logic to avoid double detections
            dets, frame = self.overlapping_removal(frame, dets)
            for i in dets:
                class_name.append(i["class"])
                # frame = cv2.rectangle(frame, (i["points"][0], i["points"][1]), (i["points"][2], i["points"][3]),
                #                       (255, 255, 0), 2)
                bboxs.append([i["points"][1], i["points"][0], i["points"][3], i["points"][2]])

        return bboxs, frame, dets, class_name

    def _predict(self, obj):
        try:
            frame = obj['frame']
            logger.info(f"Frame Received with shape {frame.shape}")
            dets, frame, _dets, class_name = self.inference(frame)
            if len(dets) > 0:
                frame, objects, boxs = self.kalman_tracker(dets, frame)
                frame = self.update_bag_count(frame=frame, detection_objects=objects, class_name=class_name,
                                              detections=_dets)
                logger.debug("self.uncounted_objects --> {}".format(self.uncounted_objects))
            frame = self.draw_line_over_image(frame)
            obj['frame'] = cv2.resize(frame, (self.config.get('FRAME_WIDTH'), self.config.
                                              get('FRAME_HEIGHT')))

        except Exception as e:
            logger.exception(f"Error: {e}", exc_info=True)
            obj['frame'] = cv2.resize(obj['frame'], (self.config.get('FRAME_WIDTH'), self.config.get('FRAME_HEIGHT')))

        return obj
