import os
from pathlib import Path
import numpy as np
import cv2

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.utils.visualizer import ColorMode



PATH_PROJECT = Path(__file__).parent.parent

def get_model():
    
    """
    This function is for the model of the project
    """
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set threshold for this model
    cfg.MODEL.WEIGHTS = str(PATH_PROJECT/"output"/"model_final.pth")  # Let training initialize from model zoo
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
    predictor = DefaultPredictor(cfg)
    return predictor


def predict_image(img_pil):
    """
    This function is for the prediction of the model
    return the image with the prediction and the areas of the objects
    """
    predictor = get_model()
    img_array = np.array(img_pil)
    outputs = predictor(img_array)
    v = Visualizer(img_array,
                   scale=1, 
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    image_output = v.get_image()[:, :, ::-1].copy()
    
    masks = outputs["instances"].pred_masks.cpu().numpy()
    class_ids = outputs["instances"].pred_classes.cpu().numpy()
    areas = [np.sum(mask) for mask in masks]
    # Add text labels with the object IDs and areas
    for i, (mask, class_id, area) in enumerate(zip(masks, class_ids, areas)):
        text = f"The id is {i}"
        text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=1)
        pos = (np.unravel_index(np.argmax(mask), mask.shape))[::-1]
        pos = (pos[0] - text_size[0]//2, pos[1] - text_size[1]//2)
        cv2.putText(image_output, text, pos, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0,0,255), thickness=1)
        
    values = {"image":image_output, "areas":areas, "masks":masks}
    return values