# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import cv2
import torch
import numpy as np
import torchvision
from .utils import convert_to_numpy


class GDINOAnnotator:
    def __init__(self, cfg, device=None):
        try:
            from groundingdino.util.inference import Model, load_model, load_image, predict
        except:
            import warnings
            warnings.warn("please pip install groundingdino package, or you can refer to models/VACE-Annotators/gdino/groundingdino-0.1.0-cp310-cp310-linux_x86_64.whl")

        grounding_dino_config_path = cfg['CONFIG_PATH']
        grounding_dino_checkpoint_path = cfg['PRETRAINED_MODEL']
        grounding_dino_tokenizer_path = cfg['TOKENIZER_PATH']  # TODO
        self.box_threshold = cfg.get('BOX_THRESHOLD', 0.25)
        self.text_threshold = cfg.get('TEXT_THRESHOLD', 0.2)
        self.iou_threshold = cfg.get('IOU_THRESHOLD', 0.5)
        self.use_nms = cfg.get('USE_NMS', True)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
        self.model = Model(model_config_path=grounding_dino_config_path,
                           model_checkpoint_path=grounding_dino_checkpoint_path,
                           device=self.device)

    def forward(self, image, classes=None, caption=None):
        image_bgr = convert_to_numpy(image)[..., ::-1]  # bgr

        if classes is not None:
            classes = [classes] if isinstance(classes, str) else classes
            detections = self.model.predict_with_classes(
                image=image_bgr,
                classes=classes,
                box_threshold=self.box_threshold,
                text_threshold=self.text_threshold
            )
        elif caption is not None:
            detections, phrases = self.model.predict_with_caption(
                image=image_bgr,
                caption=caption,
                box_threshold=self.box_threshold,
                text_threshold=self.text_threshold
            )
        else:
            raise NotImplementedError()

        if self.use_nms:
            nms_idx = torchvision.ops.nms(
                torch.from_numpy(detections.xyxy),
                torch.from_numpy(detections.confidence),
                self.iou_threshold
            ).numpy().tolist()
            detections.xyxy = detections.xyxy[nms_idx]
            detections.confidence = detections.confidence[nms_idx]
            detections.class_id = detections.class_id[nms_idx] if detections.class_id is not None else None

        boxes = detections.xyxy
        confidences = detections.confidence
        class_ids = detections.class_id
        class_names = [classes[_id] for _id in class_ids] if classes is not None else phrases

        ret_data = {
            "boxes": boxes.tolist() if boxes is not None else None,
            "confidences": confidences.tolist() if confidences is not None else None,
            "class_ids": class_ids.tolist() if class_ids is not None else None,
            "class_names": class_names if class_names is not None else None,
        }
        return ret_data


class GDINORAMAnnotator:
    def __init__(self, cfg, device=None):
        from .ram import RAMAnnotator
        from .gdino import GDINOAnnotator
        self.ram_model = RAMAnnotator(cfg['RAM'], device=device)
        self.gdino_model = GDINOAnnotator(cfg['GDINO'], device=device)

    def forward(self, image):
        ram_res = self.ram_model.forward(image)
        classes = ram_res['tag_e'] if isinstance(ram_res, dict) else ram_res
        gdino_res = self.gdino_model.forward(image, classes=classes)
        return gdino_res