import gradio as gr

from ultralytics import YOLO
import torch, torchvision
import cv2
import os
import pickle

import numpy as np

class jeysonHandler:
    def __init__(self, delta_threshold, pos_threshold, last_pos_num=5, wait_time=60, num_types=4):
        self.ids={}
        self.missing_ids=set()
        self.translator={}
        self.dt = delta_threshold
        self.pt = pos_threshold
        self.lpn = last_pos_num
        self.wt = wait_time
        self.num_types = num_types
    def reg_id(self, id):
        self.ids[id] ={
        'type_statistics':[0]*self.num_types,
        'last_positions':[],
        'last_delta':np.array((0,0)),
        'events':{
            '0':[],
            '1':[]
        }
        }
    def reg_data(self, timestamp, id, type, pos):
        pos = np.array(pos)
        self.missing_ids.discard(id)
        self.ids[id]['type_statistics'][type]+=1
        self.ids[id]['last_positions'].append((timestamp, pos))
        if len(self.ids[id]['last_positions']) > self.lpn:
            self.ids[id]['last_positions'] = self.ids[id]['last_positions'][-self.lpn:]
        deltaS = self.ids[id]['last_positions'][-1][1]-self.ids[id]['last_positions'][0][1]
        deltaT = self.ids[id]['last_positions'][-1][0]-self.ids[id]['last_positions'][0][0]
        delta = deltaS/deltaT
        self.ids[id]['last_delta'] = delta
        event = {
        'start_time': timestamp,
        'end_time': timestamp
        }
        if len(self.ids[id]['events']['0']) ==0:
            self.ids[id]['events']['0'].append(event)
        elif timestamp - self.ids[id]['events']['0'][-1]['end_time'] > self.wt:
            if self.ids[id]['events']['0'][-1]['end_time'] - self.ids[id]['events']['0'][-1]['start_time'] < self.wt/3:
                self.ids[id]['events']['0'].pop()
            self.ids[id]['events']['0'].append(event)
        else:
            self.ids[id]['events']['0'][-1]['end_time'] = timestamp
        
        if np.sqrt(np.sum(np.power(delta, 2)))> self.dt:
            if len(self.ids[id]['events']['1']) ==0:
                self.ids[id]['events']['1'].append(event)
            elif timestamp - self.ids[id]['events']['1'][-1]['end_time'] > self.wt:
                if self.ids[id]['events']['1'][-1]['end_time'] - self.ids[id]['events']['1'][-1]['start_time'] < self.wt/3:
                    self.ids[id]['events']['1'].pop()
                self.ids[id]['events']['1'].append(event)
            else:
                self.ids[id]['events']['1'][-1]['end_time'] = timestamp
    
    def add_data(self, timestamp, data):
        lost = set(self.ids.keys())
        lost = lost.difference(self.missing_ids)
        newids = []
        for id, type, pos in data:
            if id in self.translator:
                id = self.translator[id]
            if id in self.ids:
                lost.discard(id)
                self.reg_data(timestamp, id, type, pos)
            else:
                newids.append((id, type, pos))
        for id, type, pos in newids:
            distances = []
            for lostid in lost:
                last_time = self.ids[lostid]['last_positions'][-1][0]
                if timestamp - last_time > self.wt:
                    self.missing_ids.add(lostid)
                    continue
                last_pos = self.ids[lostid]['last_positions'][-1][1]
                last_delta = self.ids[lostid]['last_delta']
                predicted_pos = last_pos + last_delta*(timestamp - last_time)
                distance = np.sqrt(np.sum(np.power(pos-predicted_pos,2)))
                distances.append((distance, lostid))
            if len(distances)>0:
                mindist = min(distances)
                if mindist[0]<self.pt:
                    self.translator[id] = lostid
                    self.reg_data(timestamp, lostid, type, pos)
                else:
                    self.reg_id(id)
                    self.reg_data(timestamp, id, type, pos)
            else:
                self.reg_id(id)
                self.reg_data(timestamp, id, type, pos)
    def get_json(self, type_names):
          ajson = {}
          for id in self.ids:
              idinfo = {
              'type': type_names[np.argmax(np.array(self.ids[id]['type_statistics']))],
              'events': self.ids[id]['events']
              }
              ajson[str(id)] = idinfo
          return ajson
i = 0
def timestamp():
    global i
    i += 1
    return (i - 1) / 5

def video_model(video):
    model = YOLO('last.pt')
    predict = model.track(source=video, show=False, conf=.1)
    boxes = list(map(lambda x: (timestamp(), x.boxes.id.int().numpy(), x.boxes.cls.int().numpy(), (x.boxes.xywh[..., :2] + x.boxes.xywh[..., 2:] / 2).numpy()), predict))
    handler = jeysonHandler(3, 10)
    for x in boxes:
        timestamps = x[0]
        data = []
        for i in range(len(x[1])):
            data.append((x[1][i], x[2][i], x[3][i]))
        handler.add_data(timestamps, data)
    return handler.get_json(['crane', 'excavator', 'tractor', 'truck'])

demo = gr.Interface(video_model, 
                    gr.Video(), 
                    gr.JSON(), 
                cache_examples=True).queue()

if __name__ == "__main__":
    demo.launch(share=False)