import groundingdino.datasets.transforms as T
import torch
import numpy as np
from groundingdino.models import build_model
from groundingdino.util import box_ops
from groundingdino.util.inference import predict
from groundingdino.util.slconfig import SLConfig
from groundingdino.util.utils import clean_state_dict
from torchvision.transforms import ToTensor
from huggingface_hub import hf_hub_download
import time


def load_model_hugging_face(repo_id, filename, ckpt_config_filename, device='cpu'):
    cache_config_file = hf_hub_download(repo_id=repo_id, filename=ckpt_config_filename)

    args = SLConfig.fromfile(cache_config_file)
    model = build_model(args)
    args.device = device

    cache_file = hf_hub_download(repo_id=repo_id, filename=filename)
    checkpoint = torch.load(cache_file, map_location='cpu')
    log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)
    print(f"Model loaded from {cache_file} \n => {log}")
    model.eval()
    return model


class LangEfficientSAM:
    def __init__(self, device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
        self.device = device
        print("Device:", self.device)
        if self.device == torch.device("cpu"):
            self.sam_efficient = torch.jit.load('./models/efficientsam_s_cpu.jit')
        else:
            self.sam_efficient = torch.jit.load('./models/efficientsam_s_gpu.jit')
        ckpt_repo_id = "ShilongLiu/GroundingDINO"
        ckpt_filename = "groundingdino_swinb_cogcoor.pth"
        ckpt_config_filename = "GroundingDINO_SwinB.cfg.py"
        self.groundingdino = load_model_hugging_face(ckpt_repo_id,
                                                     ckpt_filename,
                                                     ckpt_config_filename,
                                                     self.device)

    def predict_dino(self, image_pil, text_prompt, box_threshold, text_threshold):
        start = time.time()
        transform = T.Compose([
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        image_transformed, _ = transform(image_pil, None)

        boxes, logits, phrases = predict(model=self.groundingdino,
                                         image=image_transformed,
                                         caption=text_prompt,
                                         box_threshold=box_threshold,
                                         text_threshold=text_threshold,
                                         device=self.device)
        W, H = image_pil.size
        boxes = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H])
        # print("DINO time: ", time.time() - start)

        return boxes, logits, phrases

    def predict_sam(self, image, box):
        start = time.time()
        img_tensor = ToTensor()(image).to(device=self.device)
        bbox = torch.reshape(box.clone().detach(), [1, 1, 2, 2]).to(device=self.device)
        bbox_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2]).to(device=self.device)

        predicted_logits, predicted_iou = self.sam_efficient(
            img_tensor[None, ...],
            bbox,
            bbox_labels,
        )
        predicted_logits = predicted_logits.cpu()
        all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
        predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()

        max_predicted_iou = -1
        selected_mask_using_predicted_iou = None
        for m in range(all_masks.shape[0]):
            curr_predicted_iou = predicted_iou[m]
            if (
                    curr_predicted_iou > max_predicted_iou
                    or selected_mask_using_predicted_iou is None
            ):
                max_predicted_iou = curr_predicted_iou
                selected_mask_using_predicted_iou = all_masks[m]

        # print("SAM time: ", time.time() - start)
        return selected_mask_using_predicted_iou

    def predict(self, image_pil, text_prompt, box_threshold=0.3, text_threshold=0.25):
        boxes, logits, phrases = self.predict_dino(image_pil, text_prompt, box_threshold, text_threshold)
        # masks = torch.tensor([])
        masks = []
        if len(boxes) > 0:
            for box in boxes:
                mask = self.predict_sam(image_pil, box)
                masks.append(mask)

        masks = np.array(masks)
        masks = torch.from_numpy(masks)

        return masks, boxes, phrases, logits