import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
from constants import constants
from tools import generate_detections as gdet
from tools import chart
from deep_sort.tracker import Tracker
from deep_sort.detection import Detection
from deep_sort import preprocessing, nn_matching
from tensorflow.compat.v1 import InteractiveSession
from tensorflow.compat.v1 import ConfigProto
from PIL import Image
from core.config import cfg
from tensorflow.python.saved_model import tag_constants
from core.yolov4 import filter_boxes
from absl.flags import FLAGS
from absl import app, flags, logging
import core.utils as utils
import tensorflow as tf
import time
import datetime
import matplotlib.pyplot as plt
import numpy as np
import cv2
from tools import functions as func

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

flags.DEFINE_string('weights', './checkpoints/yolov4-416',
                    'path to weights file')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_string('video', './data/video/test.mp4',
                    'path to input video or set to 0 for webcam')
flags.DEFINE_string('output', None, 'path to output video')
flags.DEFINE_string('output_format', 'XVID',
                    'codec used in VideoWriter when saving video to file')
flags.DEFINE_float('iou', 0.45, 'iou threshold')
flags.DEFINE_float('score', 0.50, 'score threshold')
flags.DEFINE_boolean('dont_show', False, 'dont show video output')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_string('name', None, 'name of crossroad')
flags.DEFINE_string('direction', None, 'direction of crossroad')

def main(_argv):
    max_cosine_distance = 0.4
    nn_budget = None
    nms_max_overlap = 1.0

    model_filename = 'model_data/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename, batch_size=1)
    metric = nn_matching.NearestNeighborDistanceMetric("cosine",
                                                       max_cosine_distance,
                                                       nn_budget)
    tracker = Tracker(metric)

    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
    input_size = FLAGS.size
    video_path = FLAGS.video

    saved_model_loaded = tf.saved_model.load(FLAGS.weights,
                                             tags=[tag_constants.SERVING])
    infer = saved_model_loaded.signatures['serving_default']

    try:
        vid = cv2.VideoCapture(int(video_path))
    except:
        vid = cv2.VideoCapture(video_path)

    out = None

    if FLAGS.output:
        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vid.get(cv2.CAP_PROP_FPS))
        codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
        out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))

    while True:

        return_value, frame = vid.read()
        if return_value:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(frame)
        else:
            print('Done!')
            break

        frame_size = frame.shape[:2]
        image_data = cv2.resize(frame, (input_size, input_size))
        image_data = image_data / 255.
        image_data = image_data[np.newaxis, ...].astype(np.float32)
        start_time = time.time()

        batch_data = tf.constant(image_data)
        pred_bbox = infer(batch_data)
        for key, value in pred_bbox.items():
            boxes = value[:, :, 0:4]
            pred_conf = value[:, :, 4:]

        boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
            boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
            scores=tf.reshape(
                pred_conf,
                (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
            max_output_size_per_class=50,
            max_total_size=50,
            iou_threshold=FLAGS.iou,
            score_threshold=FLAGS.score)

        num_objects = valid_detections.numpy()[0]
        bboxes = boxes.numpy()[0]
        bboxes = bboxes[0:int(num_objects)]
        scores = scores.numpy()[0]
        scores = scores[0:int(num_objects)]
        classes = classes.numpy()[0]
        classes = classes[0:int(num_objects)]

        original_h, original_w, _ = frame.shape
        bboxes = utils.format_boxes(bboxes, original_h, original_w)
        pred_bbox = [bboxes, scores, classes, num_objects]
        class_names = utils.read_class_names(cfg.YOLO.CLASSES)

        allowed_classes = ['car', 'bicycle', 'motorbike', 'bus', 'truck']

        names = []
        deleted_indx = []
        for i in range(num_objects):
            class_indx = int(classes[i])
            class_name = class_names[class_indx]
            if class_name not in allowed_classes:
                deleted_indx.append(i)
            else:
                names.append(class_name)
        names = np.array(names)
        count = len(names)

        bboxes = np.delete(bboxes, deleted_indx, axis=0)
        scores = np.delete(scores, deleted_indx, axis=0)

        # encode yolo detections and feed to tracker
        features = encoder(frame, bboxes)
        detections = [
            Detection(bbox, score, class_name, feature)
            for bbox, score, class_name, feature in zip(
                bboxes, scores, names, features)
        ]

        # color map
        cmap = plt.get_cmap('tab20b')
        colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

        # non-maxima suppression
        boxs = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        classes = np.array([d.class_name for d in detections])
        indices = preprocessing.non_max_suppression(boxs, classes,
                                                    nms_max_overlap, scores)
        detections = [detections[i] for i in indices]

        # tracker
        tracker.predict()
        tracker.update(detections)

        for track in tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            bbox = track.to_tlbr()
            class_name = track.get_class()

            # bbox
            rects = []
            box = np.array(
                [int(bbox[0]),
                 int(bbox[1]),
                 int(bbox[2]),
                 int(bbox[3])])
            rects.append(box.astype(int))
            centroid = np.zeros((len(rects), 2), dtype="int")
            cX = int((int(bbox[0]) + int(bbox[2])) / 2.0)
            cY = int((int(bbox[1]) + int(bbox[3])) / 2.0)
            centroid = (cX, cY)
            cv2.circle(frame, (centroid[0], centroid[1]), 4, (0, 255, 0), -1)
            color = colors[int(track.track_id) % len(colors)]
            color = [i * 255 for i in color]
            cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])),
                          (int(bbox[2]), int(bbox[3])), (color), 2)
            cv2.putText(frame, class_name + " " + str(track.track_id),
                        (int(bbox[0]), int(bbox[1] - 10)), 0, 0.4,
                        (255, 255, 255), 1)

            if track.track_id in constants.MovingObjects.keys():
                directionx = centroid[0] - \
                    constants.MovingObjects[track.track_id][0][0]
                directiony = centroid[1] - \
                    constants.MovingObjects[track.track_id][0][1]

                if constants.MovingObjects[track.track_id][1] == "":
                    entry_direction = func.get_entry_direction(
                        constants.MovingObjects[track.track_id][0][0],
                        constants.MovingObjects[track.track_id][0][1], width,
                        height)
                    entry_time = datetime.datetime.now()
                    constants.MovingObjects[
                        track.track_id][1] = entry_direction
                    constants.MovingObjects[
                        track.track_id][3] = entry_time.strftime("%X %x")

                if np.abs(directionx) > 5 and np.abs(directiony) > 5:
                    exit_direction = func.get_direction(directionx, directiony)
                    exit_time = datetime.datetime.now()
                    constants.MovingObjects[track.track_id][2] = exit_direction
                    constants.MovingObjects[
                        track.track_id][4] = exit_time.strftime("%X %x")

            else:
                constants.MovingObjects[track.track_id] = [
                    centroid, "", "", "", ""
                ]

        #compass
        direction = FLAGS.direction
        result = func.compass(frame,direction)
     
        result = np.asarray(frame)
        result = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

        if not FLAGS.dont_show:
            cv2.imshow("Output Video", result)
            
        if FLAGS.output:
            out.write(result)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            cv2.destroyAllWindows()

    #graph
    name = FLAGS.name
    chart.draw_chart(func.counters(name))
    #time text
    func.get_time()


if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
from constants import constants
from tools import generate_detections as gdet
from tools import chart
from deep_sort.tracker import Tracker
from deep_sort.detection import Detection
from deep_sort import preprocessing, nn_matching
from tensorflow.compat.v1 import InteractiveSession
from tensorflow.compat.v1 import ConfigProto
from PIL import Image
from core.config import cfg
from tensorflow.python.saved_model import tag_constants
from core.yolov4 import filter_boxes
from absl.flags import FLAGS
from absl import app, flags, logging
import core.utils as utils
import tensorflow as tf
import time
import datetime
import matplotlib.pyplot as plt
import numpy as np
import cv2
from tools import functions as func

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

flags.DEFINE_string('weights', './checkpoints/yolov4-416',
                    'path to weights file')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_string('video', './data/video/test.mp4',
                    'path to input video or set to 0 for webcam')
flags.DEFINE_string('output', None, 'path to output video')
flags.DEFINE_string('output_format', 'XVID',
                    'codec used in VideoWriter when saving video to file')
flags.DEFINE_float('iou', 0.45, 'iou threshold')
flags.DEFINE_float('score', 0.50, 'score threshold')
flags.DEFINE_boolean('dont_show', False, 'dont show video output')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_string('name', None, 'name of crossroad')
flags.DEFINE_string('direction', None, 'direction of crossroad')

def main(_argv):
    max_cosine_distance = 0.4
    nn_budget = None
    nms_max_overlap = 1.0

    model_filename = 'model_data/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename, batch_size=1)
    metric = nn_matching.NearestNeighborDistanceMetric("cosine",
                                                       max_cosine_distance,
                                                       nn_budget)
    tracker = Tracker(metric)

    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
    input_size = FLAGS.size
    video_path = FLAGS.video

    saved_model_loaded = tf.saved_model.load(FLAGS.weights,
                                             tags=[tag_constants.SERVING])
    infer = saved_model_loaded.signatures['serving_default']

    try:
        vid = cv2.VideoCapture(int(video_path))
    except:
        vid = cv2.VideoCapture(video_path)

    out = None

    if FLAGS.output:
        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vid.get(cv2.CAP_PROP_FPS))
        codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
        out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))

    while True:

        return_value, frame = vid.read()
        if return_value:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(frame)
        else:
            print('Done!')
            break

        frame_size = frame.shape[:2]
        image_data = cv2.resize(frame, (input_size, input_size))
        image_data = image_data / 255.
        image_data = image_data[np.newaxis, ...].astype(np.float32)
        start_time = time.time()

        batch_data = tf.constant(image_data)
        pred_bbox = infer(batch_data)
        for key, value in pred_bbox.items():
            boxes = value[:, :, 0:4]
            pred_conf = value[:, :, 4:]

        boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
            boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
            scores=tf.reshape(
                pred_conf,
                (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
            max_output_size_per_class=50,
            max_total_size=50,
            iou_threshold=FLAGS.iou,
            score_threshold=FLAGS.score)

        num_objects = valid_detections.numpy()[0]
        bboxes = boxes.numpy()[0]
        bboxes = bboxes[0:int(num_objects)]
        scores = scores.numpy()[0]
        scores = scores[0:int(num_objects)]
        classes = classes.numpy()[0]
        classes = classes[0:int(num_objects)]

        original_h, original_w, _ = frame.shape
        bboxes = utils.format_boxes(bboxes, original_h, original_w)
        pred_bbox = [bboxes, scores, classes, num_objects]
        class_names = utils.read_class_names(cfg.YOLO.CLASSES)

        allowed_classes = ['car', 'bicycle', 'motorbike', 'bus', 'truck']

        names = []
        deleted_indx = []
        for i in range(num_objects):
            class_indx = int(classes[i])
            class_name = class_names[class_indx]
            if class_name not in allowed_classes:
                deleted_indx.append(i)
            else:
                names.append(class_name)
        names = np.array(names)
        count = len(names)

        bboxes = np.delete(bboxes, deleted_indx, axis=0)
        scores = np.delete(scores, deleted_indx, axis=0)

        # encode yolo detections and feed to tracker
        features = encoder(frame, bboxes)
        detections = [
            Detection(bbox, score, class_name, feature)
            for bbox, score, class_name, feature in zip(
                bboxes, scores, names, features)
        ]

        # color map
        cmap = plt.get_cmap('tab20b')
        colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

        # non-maxima suppression
        boxs = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        classes = np.array([d.class_name for d in detections])
        indices = preprocessing.non_max_suppression(boxs, classes,
                                                    nms_max_overlap, scores)
        detections = [detections[i] for i in indices]

        # tracker
        tracker.predict()
        tracker.update(detections)

        for track in tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            bbox = track.to_tlbr()
            class_name = track.get_class()

            # bbox
            rects = []
            box = np.array(
                [int(bbox[0]),
                 int(bbox[1]),
                 int(bbox[2]),
                 int(bbox[3])])
            rects.append(box.astype(int))
            centroid = np.zeros((len(rects), 2), dtype="int")
            cX = int((int(bbox[0]) + int(bbox[2])) / 2.0)
            cY = int((int(bbox[1]) + int(bbox[3])) / 2.0)
            centroid = (cX, cY)
            cv2.circle(frame, (centroid[0], centroid[1]), 4, (0, 255, 0), -1)
            color = colors[int(track.track_id) % len(colors)]
            color = [i * 255 for i in color]
            cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])),
                          (int(bbox[2]), int(bbox[3])), (color), 2)
            cv2.putText(frame, class_name + " " + str(track.track_id),
                        (int(bbox[0]), int(bbox[1] - 10)), 0, 0.4,
                        (255, 255, 255), 1)

            if track.track_id in constants.MovingObjects.keys():
                directionx = centroid[0] - \
                    constants.MovingObjects[track.track_id][0][0]
                directiony = centroid[1] - \
                    constants.MovingObjects[track.track_id][0][1]

                if constants.MovingObjects[track.track_id][1] == "":
                    entry_direction = func.get_entry_direction(
                        constants.MovingObjects[track.track_id][0][0],
                        constants.MovingObjects[track.track_id][0][1], width,
                        height)
                    entry_time = datetime.datetime.now()
                    constants.MovingObjects[
                        track.track_id][1] = entry_direction
                    constants.MovingObjects[
                        track.track_id][3] = entry_time.strftime("%X %x")

                if np.abs(directionx) > 5 and np.abs(directiony) > 5:
                    exit_direction = func.get_direction(directionx, directiony)
                    exit_time = datetime.datetime.now()
                    constants.MovingObjects[track.track_id][2] = exit_direction
                    constants.MovingObjects[
                        track.track_id][4] = exit_time.strftime("%X %x")

            else:
                constants.MovingObjects[track.track_id] = [
                    centroid, "", "", "", ""
                ]

        #compass
        direction = FLAGS.direction
        result = func.compass(frame,direction)
     
        result = np.asarray(frame)
        result = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

        if not FLAGS.dont_show:
            cv2.imshow("Output Video", result)
            
        if FLAGS.output:
            out.write(result)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            cv2.destroyAllWindows()

    #graph
    name = FLAGS.name
    chart.draw_chart(func.counters(name))
    #time text
    func.get_time()


if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass

# encode yolo detections and feed to tracker # color map # non-maxima suppression # tracker # bbox #compass #graph #time text