# layerstyle advance
import folder_paths
from .imagefunc import *

select_list = ["all", "first", "by_index"]
sort_method_list = ["left_to_right", "top_to_bottom", "big_to_small", "confidence"]


def sort_bboxes(bboxes:list, method:str) -> list:
    sorted_bboxes = []
    if method == "left_to_right":
        sorted_bboxes = sorted(bboxes, key=lambda bbox: bbox[0])
    elif method == "top_to_bottom":
        sorted_bboxes = sorted(bboxes, key=lambda bbox: bbox[1])
    elif method == "big_to_small":
        sorted_bboxes = sorted(bboxes, key=lambda bbox: (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), reverse=True)
    else:
        sorted_bboxes = bboxes
    return sorted_bboxes

def select_bboxes(bboxes:list, bbox_select:str, select_index:str) -> list:
    indexs = extract_numbers(select_index)
    if bbox_select == "all":
        return bboxes
    elif bbox_select == "first":
        return [bboxes[0]]
    elif bbox_select == "by_index":
        new_bboxes = []
        for i in indexs:
            try:
                new_bboxes.append(bboxes[i])
            except IndexError:
                log(f"Object detector output by_index: invalid bbox index {i}", message_type='warning')
        return new_bboxes


class LS_BBOXES_JOIN:

    def __init__(self):
        self.NODE_NAME = 'BBoxes Join'

    @classmethod
    def INPUT_TYPES(cls):

        return {
            "required": {
                "bboxes_1": ("BBOXES",),
            },
            "optional": {
                "bboxes_2": ("BBOXES",),
                "bboxes_3": ("BBOXES",),
                "bboxes_4": ("BBOXES",),
            }
        }

    RETURN_TYPES = ("BBOXES",)
    RETURN_NAMES = ("bboxes",)
    FUNCTION = 'bboxes_join'
    CATEGORY = '😺dzNodes/LayerMask'

    def bboxes_join(self, bboxes_1, bboxes_2=None, bboxes_3=None, bboxes_4=None):
        all_inputs = [b for b in [bboxes_2, bboxes_3, bboxes_4] if b is not None]
        for other in all_inputs:
            for i in range(len(other)):
                if i < len(bboxes_1):
                    bboxes_1[i].extend(other[i])
                else:
                    bboxes_1.append(other[i])
        return (bboxes_1,)

class LS_OBJECT_DETECTOR_FL2:

    def __init__(self):
        self.NODE_NAME = 'Object Detector Florence2'

    @classmethod
    def INPUT_TYPES(cls):

        return {
            "required": {
                "image": ("IMAGE", ),  #
                "prompt": ("STRING", {"default": "subject"}),
                "florence2_model": ("FLORENCE2",),
                "sort_method": (sort_method_list,),
                "bbox_select": (select_list,),
                "select_index": ("STRING", {"default": "0,"},),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("BBOXES", "IMAGE",)
    RETURN_NAMES = ("bboxes", "preview",)
    FUNCTION = 'object_detector_fl2'
    CATEGORY = '😺dzNodes/LayerMask'

    def object_detector_fl2(self, image, prompt, florence2_model, sort_method, bbox_select, select_index):

        ret_bboxes = []
        ret_previews = []
        max_new_tokens = 512
        num_beams = 3
        do_sample = False
        fill_mask = False

        model = florence2_model['model']
        processor = florence2_model['processor']

        for img in image:
            bboxes = []
            img = tensor2pil(img.unsqueeze(0)).convert("RGB")
            task = 'caption to phrase grounding'
            from .florence2_ultra import  process_image
            results, _ = process_image(model, processor, img, task,
                                       max_new_tokens, num_beams, do_sample,
                                       fill_mask, prompt)

            if isinstance(results, dict):
                results["width"] = img.width
                results["height"] = img.height

            bboxes = self.fbboxes_to_list(results)
            bboxes = sort_bboxes(bboxes, sort_method)
            bboxes = select_bboxes(bboxes, bbox_select, select_index)
            preview = draw_bounding_boxes(img, bboxes, color="random", line_width=-1)
            ret_previews.append(pil2tensor(preview))
            ret_bboxes.append(standardize_bbox(bboxes))

            # if len(bboxes) == 0:
            #     log(f"{self.NODE_NAME} no object found", message_type='warning')
            # else:
            #     log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')

        return (ret_bboxes, torch.cat(ret_previews, dim=0))

    def fbboxes_to_list(self, F_BBOXES) -> list:
        if isinstance(F_BBOXES, str):
            return None
        ret_bboxes = []
        width = F_BBOXES["width"]
        height = F_BBOXES["height"]
        x1_c = width
        y1_c = height
        x2_c = y2_c = 0
        label = ""
        if "bboxes" in F_BBOXES:
            for idx in range(len(F_BBOXES["bboxes"])):
                bbox = F_BBOXES["bboxes"][idx]
                new_label = F_BBOXES["labels"][idx].removeprefix("</s>")
                if new_label not in label:
                    if idx > 0:
                        label = label + ", "
                    label = label + new_label
                if len(bbox) == 4:
                    x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
                elif len(bbox) == 8:
                    x1 = int(min(bbox[0::2]))
                    x2 = int(max(bbox[0::2]))
                    y1 = int(min(bbox[1::2]))
                    y2 = int(max(bbox[1::2]))
                else:
                    continue
                x1_c = min(x1_c, x1)
                y1_c = min(y1_c, y1)
                x2_c = max(x2_c, x2)
                y2_c = max(y2_c, y2)
                ret_bboxes.append([x1, y1, x2, y2])
        else:
            x1_c = width
            y1_c = height
            x2_c = y2_c = 0
            for polygon in F_BBOXES["polygons"][0]:
                if len(polygon) < 3:
                    print('Invalid polygon:', polygon)
                    continue
                x1_c = min(x1_c, int(min(polygon[0::2])))
                x2_c = max(x2_c, int(max(polygon[0::2])))
                y1_c = min(y1_c, int(min(polygon[1::2])))
                y2_c = max(y2_c, int(max(polygon[1::2])))
            ret_bboxes.append([x1_c, y1_c, x2_c, y2_c])
        if len(ret_bboxes) == 0:
            ret_bboxes.append([x1_c, y1_c, x2_c, y2_c])
        return ret_bboxes

class LS_OBJECT_DETECTOR_MASK:

    def __init__(self):
        self.NODE_NAME = 'Object Detector MASK'

    @classmethod
    def INPUT_TYPES(cls):

        return {
            "required": {
                "object_mask": ("MASK",),
                "sort_method": (sort_method_list,),
                "bbox_select": (select_list,),
                "select_index": ("STRING", {"default": "0,"},),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("BBOXES", "IMAGE",)
    RETURN_NAMES = ("bboxes", "preview",)
    FUNCTION = 'object_detector_mask'
    CATEGORY = '😺dzNodes/LayerMask'

    def object_detector_mask(self, object_mask, sort_method, bbox_select, select_index):

        ret_bboxes = []
        ret_previews = []

        if object_mask.dim() == 2:
            object_mask = torch.unsqueeze(object_mask, 0)

        for msk in object_mask:
            bboxes = []
            cv_mask = tensor2cv2(msk)
            cv_mask = cv2.cvtColor(cv_mask, cv2.COLOR_BGR2GRAY)
            _, binary = cv2.threshold(cv_mask, 127, 255, cv2.THRESH_BINARY)
            # invert mask
            # binary = cv2.bitwise_not(binary)
            contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for contour in contours:
                x, y, w, h = cv2.boundingRect(contour)
                bboxes.append([x, y, x + w, y + h])
            bboxes = sort_bboxes(bboxes, sort_method)
            bboxes = select_bboxes(bboxes, bbox_select, select_index)
            preview = draw_bounding_boxes(tensor2pil(msk).convert("RGB"), bboxes, color="random", line_width=-1)
            ret_previews.append(pil2tensor(preview))

            # if len(bboxes) == 0:
            #     log(f"{self.NODE_NAME} no object found", message_type='warning')
            # else:
            #     log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')

            ret_bboxes.append(standardize_bbox(bboxes))

        return (ret_bboxes, torch.cat(ret_previews, dim=0))


class LS_OBJECT_DETECTOR_YOLO8:

    def __init__(self):
        self.NODE_NAME = 'Object Detector YOLO8'

    @classmethod
    def INPUT_TYPES(cls):
        model_ext = [".pt"]
        model_path = os.path.join(folder_paths.models_dir, 'yolo')
        FILES_DICT = get_files(model_path, model_ext)
        FILE_LIST = list(FILES_DICT.keys())
        return {
            "required": {
                "image": ("IMAGE", ),
                "yolo_model": (FILE_LIST,),
                "sort_method": (sort_method_list,),
                "bbox_select": (select_list,),
                "select_index": ("STRING", {"default": "0,"},),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("BBOXES", "IMAGE",)
    RETURN_NAMES = ("bboxes", "preview",)
    FUNCTION = 'object_detector_yolo8'
    CATEGORY = '😺dzNodes/LayerMask'

    def object_detector_yolo8(self, image, yolo_model, sort_method, bbox_select, select_index):

        from  ultralytics import YOLO
        model_path = os.path.join(folder_paths.models_dir, 'yolo')
        yolo_model = YOLO(os.path.join(model_path, yolo_model))

        ret_bboxes = []
        ret_previews = []

        for img in image:
            bboxes = []
            img = torch.unsqueeze(img.unsqueeze(0), 0)
            _image = tensor2pil(img)
            results = yolo_model(_image, retina_masks=True)
            for result in results:
                yolo_plot_image = cv2.cvtColor(result.plot(), cv2.COLOR_BGR2RGB)

                # no mask, if have box, draw box
                if result.boxes is not None and len(result.boxes.xyxy) > 0:
                    for box in result.boxes:
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        bboxes.append([x1, y1, x2, y2])
            bboxes = sort_bboxes(bboxes, sort_method)
            bboxes = select_bboxes(bboxes, bbox_select, select_index)
            preview = draw_bounding_boxes(_image.convert("RGB"), bboxes, color="random", line_width=-1)
            ret_previews.append(pil2tensor(preview))

            # if len(bboxes) == 0:
            #     log(f"{self.NODE_NAME} no object found", message_type='warning')
            # else:
            #     log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')

            ret_bboxes.append(standardize_bbox(bboxes))

        return (ret_bboxes, torch.cat(ret_previews, dim=0),)

class LS_OBJECT_DETECTOR_YOLOWORLD:

    def __init__(self):
        self.NODE_NAME = 'Object Detector YOLO-WORLD'
        self.model_path = os.path.join(folder_paths.models_dir, 'yolo-world')
        os.environ['MODEL_CACHE_DIR'] = self.model_path

    @classmethod
    def INPUT_TYPES(cls):
        model_list =['yolo_world/v2-x', 'yolo_world/v2-l', 'yolo_world/v2-m',
                    'yolo_world/v2-s', 'yolo_world/l', 'yolo_world/m',
                    'yolo_world/s']
        return {
            "required": {
                "image": ("IMAGE", ),
                "yolo_world_model": (model_list,),
                "confidence_threshold": ("FLOAT", {"default": 0.05, "min": 0, "max": 1, "step": 0.01}),
                "nms_iou_threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1, "step": 0.01}),
                "prompt": ("STRING", {"default": "subject"}),
                "sort_method": (sort_method_list,),
                "bbox_select": (select_list,),
                "select_index": ("STRING", {"default": "0,"},),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("BBOXES", "IMAGE",)
    RETURN_NAMES = ("bboxes", "preview",)
    FUNCTION = 'object_detector_yoloworld'
    CATEGORY = '😺dzNodes/LayerMask'

    def object_detector_yoloworld(self, image, yolo_world_model,
                                  confidence_threshold, nms_iou_threshold, prompt,
                                  sort_method, bbox_select, select_index):
        ret_previews = []
        ret_bboxes = []

        try:
            import supervision as sv
        except ImportError as e:
            log(f"{self.NODE_NAME}: {e}", message_type='warning')
            return None
        model=self.load_yolo_world_model(yolo_world_model, prompt)

        for i in image:
            infer_outputs = []
            # img = (255 * img.unsqueeze(0).cpu().numpy()).astype(np.uint8)
            img = tensor2np(i)
            results = model.infer(
                img, confidence=confidence_threshold)
            detections = sv.Detections.from_inference(results)
            detections = detections.with_nms(
                class_agnostic=False,
                threshold=nms_iou_threshold
            )
            infer_outputs.append(detections)

            # if len(infer_outputs[0].xyxy) > 0:
            #     bboxes = infer_outputs[0].xyxy.tolist()
            #     bboxes = [[int(value) for value in sublist] for sublist in bboxes]
            #     bboxes = sort_bboxes(bboxes, sort_method)
            #     bboxes = select_bboxes(bboxes, bbox_select, select_index)
            # else:
            #     bboxes = []

            bboxes = infer_outputs[0].xyxy.tolist()
            bboxes = [[int(value) for value in sublist] for sublist in bboxes]
            bboxes = sort_bboxes(bboxes, sort_method)
            bboxes = select_bboxes(bboxes, bbox_select, select_index)


            preview = draw_bounding_boxes(tensor2pil(i.unsqueeze(0)).convert('RGB'), bboxes, color="random", line_width=-1)
            ret_previews.append(pil2tensor(preview))

            # if len(bboxes) == 0:
            #     log(f"{self.NODE_NAME} no object found", message_type='warning')
            # else:
            #     log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')

            ret_bboxes.append(standardize_bbox(bboxes))

        return (ret_bboxes, torch.cat(ret_previews, dim=0))

    def process_categories(self, categories: str) -> List[str]:
        return [category.strip().lower() for category in categories.split(',')]

    def load_yolo_world_model(self,model_id: str, categories: str) -> List[torch.nn.Module]:
        try:
            from inference.models import YOLOWorld as YOLOWorldImpl
        except ImportError as e:
            log(f"{self.NODE_NAME}: {e}", message_type='warning')
            return None
        model = YOLOWorldImpl(model_id=model_id)
        categories = self.process_categories(categories)
        model.set_classes(categories)
        return model


class LS_DrawBBoxMask:

    def __init__(self):
        self.NODE_NAME = 'Draw BBOX Mask'
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "bboxes": ("BBOXES",),
                "grow_top": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}), # bbox向上扩展，按高度比例
                "grow_bottom": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
                "grow_left": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
                "grow_right": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("mask",)
    FUNCTION = 'draw_bbox_mask'
    CATEGORY = '😺dzNodes/LayerMask'

    def draw_bbox_mask(self, image, bboxes, grow_top, grow_bottom, grow_left, grow_right
                      ):

        ret_masks = []
        for index in range(len(image)):
            img = tensor2pil(image[index].unsqueeze(0))
            mask = Image.new("L", img.size, color='black')
            bboxes_i = bboxes[index]
            for bbox in bboxes_i:
                try:
                    if len(bbox) == 0:
                        continue
                    else:
                        x1, y1, x2, y2 = bbox
                except ValueError:
                    if len(bbox) == 0:
                        continue
                    else:
                        x1, y1, x2, y2 = bbox[index]
                w = x2 - x1
                h = y2 - y1
                if grow_top:
                    y1 = int(y1 - h * grow_top)
                if grow_bottom:
                    y2 = int(y2 + h * grow_bottom)
                if grow_left:
                    x1 = int(x1 - w * grow_left)
                if grow_right:
                    x2 = int(x2 + w * grow_right)
                if y1 > y2 or x1 > x2:
                    log(f"{self.NODE_NAME} Invalid bbox after extend: ({x1},{y1},{x2},{y2})", message_type='warning')
                    continue
                draw = ImageDraw.Draw(mask)
                draw.rectangle([x1, y1, x2, y2], fill='white', outline='white', width=0)
            ret_masks.append(pil2tensor(mask))

        log(f"{self.NODE_NAME} Processed {len(ret_masks)} mask(s).", message_type='finish')
        return (torch.cat(ret_masks, dim=0),)


class LS_DrawBBoxMaskV2:

    def __init__(self):
        self.NODE_NAME = 'Draw BBOX Mask V2'
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "bboxes": ("BBOXES",),
                "grow_top": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}), # bbox向上扩展，按高度比例
                "grow_bottom": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
                "grow_left": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
                "grow_right": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
                "rounded_rect_radius": ("INT", {"default": 50, "min": 0, "max": 100, "step": 1}),
                "anti_aliasing": ("INT", {"default": 2, "min": 0, "max": 16, "step": 1}),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("mask",)
    FUNCTION = 'draw_bbox_mask_v2'
    CATEGORY = '😺dzNodes/LayerMask'

    def draw_bbox_mask_v2(self, image, bboxes, grow_top, grow_bottom, grow_left, grow_right,
                          rounded_rect_radius, anti_aliasing):

        ret_masks = []
        for index in range(len(image)):
            img = tensor2pil(image[index].unsqueeze(0))
            mask = Image.new("L", img.size, color='black')
            bboxes_i = bboxes[index]
            if grow_top or grow_bottom or grow_left or grow_right:
                new_bboxes_i = []
                for bbox in bboxes_i:
                    try:
                        if len(bbox) == 0:
                            continue
                        else:
                            x1, y1, x2, y2 = bbox
                    except ValueError:
                        if len(bbox) == 0:
                            continue
                        else:
                            x1, y1, x2, y2 = bbox[index]
                    w = x2 - x1
                    h = y2 - y1
                    if grow_top:
                        y1 = int(y1 - h * grow_top)
                    if grow_bottom:
                        y2 = int(y2 + h * grow_bottom)
                    if grow_left:
                        x1 = int(x1 - w * grow_left)
                    if grow_right:
                        x2 = int(x2 + w * grow_right)
                    if y1 > y2:
                        y1, y2 = y2, y1
                    if x1 > x2:
                        x1, x2 = x2, x1
                    if y2 - y1 < 1:
                        y2 += 1
                    if x2 - x1 < 1:
                        x2 += 1
                    new_bboxes_i.append((x1, y1, x2, y2))
                bboxes_i = new_bboxes_i
            mask = draw_rounded_rectangle(mask, rounded_rect_radius, bboxes_i, anti_aliasing)
            ret_masks.append(pil2tensor(mask))

        log(f"{self.NODE_NAME} Processed {len(ret_masks)} mask(s).", message_type='finish')
        return (torch.cat(ret_masks, dim=0),)


NODE_CLASS_MAPPINGS = {
    "LayerMask: BBoxJoin": LS_BBOXES_JOIN,
    "LayerMask: DrawBBoxMaskV2": LS_DrawBBoxMaskV2,
    "LayerMask: DrawBBoxMask": LS_DrawBBoxMask,
    "LayerMask: ObjectDetectorFL2": LS_OBJECT_DETECTOR_FL2,
    "LayerMask: ObjectDetectorMask": LS_OBJECT_DETECTOR_MASK,
    "LayerMask: ObjectDetectorYOLO8": LS_OBJECT_DETECTOR_YOLO8,
    "LayerMask: ObjectDetectorYOLOWorld": LS_OBJECT_DETECTOR_YOLOWORLD
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LayerMask: BBoxJoin": "LayerMask: BBox Join(Advance)",
    "LayerMask: DrawBBoxMaskV2": "LayerMask: Draw BBox Mask V2(Advance)",
    "LayerMask: DrawBBoxMask": "LayerMask: Draw BBox Mask(Advance)",
    "LayerMask: ObjectDetectorFL2": "LayerMask: Object Detector Florence2(Advance)",
    "LayerMask: ObjectDetectorMask": "LayerMask: Object Detector Mask(Advance)",
    "LayerMask: ObjectDetectorYOLO8": "LayerMask: Object Detector YOLO8(Advance)",
    "LayerMask: ObjectDetectorYOLOWorld": "LayerMask: Object Detector YOLO World(Obsolete)"
}