import torch
import cv2
import os
import sys
import numpy as np
import time
#from queue import Queue
from multiprocessing import Process, Queue
from ultralytics import YOLO
from threading import Thread
import logging
from nodes.Node import Node
from ultralytics.utils.plotting import Annotator, colors

class FishFilterNode(Node, Process):
    def __init__(self,name="FishFilterNode",parent=None):
        super().__init__(name, parent)
        super(Process,self).__init__()

        self.is_stop=False

    def run(self):
        self.setStatus(Node.NODE_STATUS_RUNNING)
        self.is_stopped = False

        nextFishId = 0
        fishMap = {}
        objLife = {}

        while True:
            if self.is_stopped:
                break
            else:
                if( not self._framePool.empty()):

                    frame, videoIndex, frameIndex, tracked_dets = self._framePool.get()

                    annotator = Annotator(frame, line_width=4)

                    if len(tracked_dets)>0:
                        bbox_xyxy = tracked_dets[:,:4]
                        identities = tracked_dets[:, 8]
                        categories = tracked_dets[:, 4]

                        for i, box in enumerate(bbox_xyxy):
                            clsIdx = int(categories[i])
                            objIdx = int(identities[i])

                            fishId = 0

                            if objIdx not in fishMap:
                                fishMap[objIdx]=nextFishId
                                fishId = nextFishId
                                nextFishId += 1
                            else:
                                fishId = fishMap[objIdx]

                            tracked_dets[i, 8]=fishId

                            objLife[objIdx]=4


                            #if( categories[i] == 0 ):
                            annotator.box_label(box.tolist(), str(fishId), color=colors(0, True))
                            #else:
                            #annotator.box_label(box.tolist(),labels[clsIdx] + ":" + str(objIdx),color=colors(clsIdx,True))

                    for oid in objLife:
                        objLife[oid]-=1
                        if objLife[oid] < 0:
                            del objLife[oid]
                            del fishMap[oid]
                            break

                    if frame is not None:
                        self.passFrame([frame, videoIndex, frameIndex, tracked_dets])
                        logging.info("FishFilter filter a frame.")


            if self._framePool.empty():
                if self.getParentStatus() == Node.NODE_STATUS_EXITED:
                    break

        self.setStatus(Node.NODE_STATUS_EXITED)
        logging.info("=========>>>>>>>>>>>>>>>>>>>>")
        logging.info("FishFilterNode is completed.")
        sys.exit()

    def stop (self):
        self.is_stop = True
        self._parent.stop()