
from time import sleep
import threading
import cv2
import numpy as np
from PIL import Image

from Object import createObjects, computeDistanceBetweenObjects, Object

from utils import mse, initProgram, computeCurveDifference, processMemory, timing, drawBoxes
from correlation import Correlation
from ProcessedImage import ProcessedImage

from yolo3.yolo import YOLO
import os

##### Initialization
resolution = (1024,1024)
inputVideo = 'images/train_station.mp4'
fourcc = cv2.VideoWriter_fourcc(*'XVID')
#out = cv2.VideoWriter(inputVideo + '_output.avi', fourcc, 20.0, resolution)
imgs = []
processedImages = []
imagesToShow = []
cropped = []
debug = 1

##### Thread for object detection. Creates array of images with objects

def processImage():


    cap = cv2.VideoCapture(inputVideo)




    yoloNetwork = YOLO()

    i = 0
    while (i < 2800):
        cap.read()
        i = i + 1

    while(cap.isOpened()):



        ret, img = cap.read()
        img = cv2.resize(np.asarray(img), resolution)
        img = Image.fromarray(img)

        out_classes, class_names, out_boxes, out_scores = yoloNetwork.getObjectsFromImage(img)


        objects = createObjects(out_classes, class_names, out_boxes)

        processedImages.append(ProcessedImage(img, objects))

    cap.release()

def objectTracking():
    memory = []

    maxXY = 80
    class_names, anchors = initProgram()

    while 1:

        if len(processedImages) > 0:

            processedImage =  processedImages[0]
            del processedImages[0]

            img = processedImage.getImage()
            newObjects = processedImage.getObjects()
            for newObject in newObjects:
                newObject.img =  np.asarray(img)[int(newObject.top):int(newObject.bottom),
                                int(newObject.left):int(newObject.right), :]
                newObject.imgResized = cv2.resize(newObject.img, (100, 100))

            if memory != None:

                for newObject in newObjects:

                    mins = []

                    for objectFromMemory in memory:
                        mins.append(mse(objectFromMemory.imgResized, newObject.imgResized) * computeDistanceBetweenObjects(objectFromMemory, newObject))


                    if len(mins) > 0:
                        if (min(mins) < 1000000):

                            if len(mins) > 0:
                                minIndex = mins.index(min(mins))
                                mins.pop(minIndex)

                            if len(memory) > 0:
                                if (newObject.classification == memory[minIndex].classification):
                                    newObject.name = memory[minIndex].name
                                    newObject.classification = memory[minIndex].classification
                                    newObject.X = memory[minIndex].X
                                    newObject.Y = memory[minIndex].Y
                                    newObject.computeCenterPosition()
                                    del memory[minIndex]


            peoples = []
            baggages= []
            correlations = []

            for newObject in newObjects:
                if(newObject.classification == "handbag" or newObject.classification == "suitcase" or newObject.classification == 'backpack'):
                    baggages.append(Object(name=newObject.name, img=newObject.img))
                    baggages[len(baggages)-1].X = newObject.X
                    baggages[len(baggages) - 1].left = newObject.left
                    baggages[len(baggages) - 1].right = newObject.right
                    baggages[len(baggages) - 1].top = newObject.top
                    baggages[len(baggages) - 1].bottom = newObject.bottom
                    baggages[len(baggages)-1].Y = newObject.Y
                    baggages[len(baggages) - 1].lastOwner = newObject.lastOwner
                    if (len(baggages[len(baggages) - 1].X) > maxXY+1):
                        del baggages[len(baggages) - 1].X[0]
                        del baggages[len(baggages) - 1].Y[0]
                else:
                    peoples.append(Object(name=newObject.name,img=newObject.img))
                    peoples[len(peoples) - 1].X = newObject.X
                    peoples[len(peoples) - 1].Y = newObject.Y
                    if (len(peoples[len(peoples) - 1].X) > maxXY+1):
                        del peoples[len(peoples) - 1].X[0]
                        del peoples[len(peoples) - 1].Y[0]

            alarms = []

            for baggage in baggages:

                if len(baggage.X) > maxXY:
                    X, Y = baggage.getDeltas()



                    mins = []
                    for person in peoples:
                        mins.append(computeCurveDifference(baggage, person))
                    if(len(mins) > 0):
                        minIndex = mins.index(min(mins))



                        if((mins[minIndex] < 200 * maxXY and mins[minIndex] != 0  and len(baggage.X) > 0 and len(peoples[minIndex].X) > 0) or peoples[minIndex].name == baggage.lastOwner):

                            x1 = baggage.X[len(baggage.X)-1]
                            y1 = baggage.Y[len(baggage.Y)-1]
                            x2 = peoples[minIndex].X[len(peoples[minIndex].X)-1]
                            y2 = peoples[minIndex].Y[len(peoples[minIndex].Y)-1]
                            baggage.owner = peoples[minIndex].name
                            correlations.append(Correlation(x1,y1,x2,y2))


                        trashholdMinMax = 4
                        if min(X) > -trashholdMinMax and min(Y) > -trashholdMinMax and max(X) < trashholdMinMax and max(
                                Y) < trashholdMinMax and baggage.owner == "":
                            alarms.append(baggage)

                        if (len(baggage.X) > maxXY):
                            del baggage.X[0]
                            del baggage.Y[0]


            imgProc = drawBoxes(img, newObjects, class_names, correlations, debug, alarms)
            imagesToShow.append(imgProc)



            processMemory(newObjects, memory)
            print(len(memory))

        sleep(.02)

##### Showing image

def showImages():

    if len(imagesToShow) > 0:
        imgProc = imagesToShow[0]
        del imagesToShow[0]

        cv2.startWindowThread()
        imgProc = np.asarray(imgProc)
        cv2.imshow('frame', imgProc)
        #out.write(imgProc)

        cv2.waitKey(0)



    #cv2.destroyAllWindows()



ImageProcessor = threading.Thread(target=processImage, args=())
ObjectTracking = threading.Thread(target=objectTracking, args=())

ImageProcessor.start()
ObjectTracking.start()

while ImageProcessor.is_alive():
    showImages()
    sleep(0.03)

#out.release()
ObjectTracking.join()




