import urllib.request
import tempfile


OPTIMAL_NMS_THRESHOLD = 0.7
_model_config_url = "https://huggingface.co/TZProject/final_tz_segmentor/resolve/main/final_model_config.yaml"
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():
    # return None
    import torch
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg

    ## define relevant parameters
    cfg = get_cfg()

    with tempfile.NamedTemporaryFile(suffix=".yaml") as tmp:
        print(tmp.name)
        urllib.request.urlretrieve(_model_config_url, filename=tmp.name)
        cfg.merge_from_file(tmp.name)
    if not torch.cuda.is_available():
        cfg.MODEL.DEVICE = "cpu"
    else:
        cfg.MODEL.DEVICE = 'cuda'

    ## when rerouting to use the final model (final_tz_segmentor) USE_FED_LOSS has to be set to False
    ## this setting requires the training data to calculate class imbalance that the app will not have access to and cause a runtime error
    ## some messages will appear when using the model that certain weights are not being used 
    ## but these are used during training and not inference and shouldn't affect the model performance
    ## code below
    cfg.MODEL.ROI_BOX_HEAD.USE_FED_LOSS = False
    
    predictor = DefaultPredictor(cfg)

    return predictor

def mask_nms(masks, scores, nms_threshold=OPTIMAL_NMS_THRESHOLD):
    """
    Runs class agnostic NMS on masks/segmentations instead of the bounding boxes.
    :param      masks: (list float) List of coordinates that make up the mask output from the model.
    :param      scores: (list float) List of corresponding confidence scores given to each mask.
    :param      nms_threshold: (float) Threshold to apply mask-based class agnostic NMS.
    :return     masks_kept (list float): List of masks kept after applying NMS.
    """
    import supervision as sv
    from shapely.geometry.polygon import Polygon
    
    polygons = []
    for mask in masks:
        contour = sv.mask_to_polygons(mask)
        if len(contour) > 0:
            polygons.append(Polygon(contour[0]))
        else:
            polygons.append(Polygon([]))
    order = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    masks_kept = []
    while order:
        i = order.pop(0)
        masks_kept.append(i)
        for j in order:
            # Calculate the IoU between the two polygons
            intersection = polygons[i].intersection(polygons[j]).area
            union = polygons[i].union(polygons[j]).area
            iou = intersection / union

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

def apply_nms(prediction, mask=False, cls_agnostic_nms=OPTIMAL_NMS_THRESHOLD):
    from torchvision.ops import nms
    from detectron2.structures import Instances
    
    if mask:
        nms_indices = mask_nms(prediction["instances"].pred_masks.numpy(),
                               prediction["instances"]._fields["scores"], cls_agnostic_nms)
    else:
        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()