import numpy as np
import time
import cv2
import demo_const as AI_CONST
import copy

from load_models import model_manager
from ultralytics.data.loaders import LoadStreams,LOGGER

def class_info():
    od_model = model_manager.get_od()  # pt파일에 따라 값이 달라짐

    class_dict = copy.deepcopy(od_model.model.names)
    class_dict[1000] = 'signalman'
    return class_dict


CLASS_INFORMATION = {}
if not CLASS_INFORMATION:
    CLASS_INFORMATION = class_info()
    CLASS_SWAP_INFO = {v:k for k,v in CLASS_INFORMATION.items()}
    
class LoadStreamsDaool(LoadStreams):
    
    def update(self, i, cap, stream):
        """Read stream `i` frames in daemon thread."""
        n, f = 0, self.frames[i]  # frame number, frame array
        while self.running and cap.isOpened() and n < (f - 1):
            if len(self.imgs[i]) < AI_CONST.LOADSTREAMS_IMG_BUFFER:  # keep a <=30-image buffer
                n += 1
                cap.grab()  # .read() = .grab() followed by .retrieve()
                if n % self.vid_stride == 0:
                    success, im = cap.retrieve()
                    if not success:
                        im = np.zeros(self.shape[i], dtype=np.uint8)
                        LOGGER.warning("WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.")
                        cap.open(stream)  # re-open stream if signal was lost
                    if self.buffer:
                        self.imgs[i].append(im)
                    else:
                        self.imgs[i] = [im]
            else:
                time.sleep(0.01)  # wait until the buffer is empty

class CustomVideoCapture:
    def __init__(self) -> None:
        self.fps = 30
        # self.fourcc = cv2.VideoWriter_fourcc(*'X264')
        self.fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.size = None
        self.video_writer = None

    def set_frame_size(self, image):
        if isinstance(image, np.ndarray):
            height,width,ch = image.shape
            self.size = (int(width), int (height))

    def set_video_writer(self, path):
        if self.size == None:
            raise
        
        self.video_writer = cv2.VideoWriter(path, self.fourcc, self.fps, self.size)
    
    def write_video(self,frame):
        if self.video_writer:
            self.video_writer.write(frame)
            
    def release_video(self):
        if self.video_writer:
            self.video_writer.release()
            self.video_writer = None


def get_monitorsize():
    width = 0
    height = 0
    try:
        import pyautogui #pip install pyautogui
        width = pyautogui.size()[0] # getting the width of the screen
        height = pyautogui.size()[1] # getting the height of the screen
    
    except Exception:
        (width,height) = AI_CONST.FHD_RESOLUTION
        
    finally:
        return (width,height)
    
def img_resize(image,size):
    return cv2.resize(image, dsize=size, interpolation=cv2.INTER_AREA)