def get_set_up():
    import torch
    TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
    CUDA_VERSION = torch.__version__.split("+")[-1]
    print("torch: ", TORCH_VERSION, "; cuda: ", CUDA_VERSION)
    print(f'GPU available: {torch.cuda.is_available()}')
    print(torch.cuda.get_device_capability())

    # print("detectron2:", detectron2.__version__)

def load_model():
    # def predictor(img):
    #     return {}
    # return predictor
    # import some common detectron2 utilities
    import torch
    from detectron2 import model_zoo
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg
    from detectron2.data.datasets import register_coco_instances

    import os
    import numpy as np

    ## define relevant parameters
    cfg = get_cfg()
    cfg.merge_from_file("./configs/test_model_config.yaml")
    if not torch.cuda.is_available():
        cfg.MODEL.DEVICE = "cpu"
    else:
        cfg.MODEL.DEVICE = 'cuda'
    predictor = DefaultPredictor(cfg)

    return predictor

def mask_nms(masks, scores, nms_threshold=0.5):
    import supervision as sv
    
    order = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    keep = []
    while order:
        i = order.pop(0)
        keep.append(i)
        for j in order:

            intersection = masks[i] * masks[j]
            union = masks[i] + masks[j]
            iou = intersection.sum() / union.sum()

            # Remove masks with IoU greater than the threshold
            if iou > nms_threshold:
                order.remove(j)
    return keep

def apply_nms(prediction, mask=False, cls_agnostic_nms=0.8):
    from torchvision.ops import nms
    from detectron2.structures import Instances
    
    if mask:
        # print("Applying mask NMS")
        nms_indices = mask_nms(prediction["instances"].pred_masks.numpy(),
                               prediction["instances"]._fields["scores"], cls_agnostic_nms)
    else:
        # print("Applying box NMS")
        nms_indices = nms(prediction["instances"].pred_boxes.tensor,
                          prediction["instances"].scores, cls_agnostic_nms)

    pred = {"instances": Instances(image_size=prediction["instances"].image_size,
                                   pred_boxes=prediction["instances"].pred_boxes[nms_indices],
                                   scores=prediction["instances"].scores[nms_indices],
                                   pred_classes=prediction["instances"].pred_classes[nms_indices],
                                   pred_masks=prediction["instances"].pred_masks[nms_indices])}

    return pred

if __name__ == '__main__':
    # get_set_up()
    load_model()