from ultralyticsplus import YOLO
from typing import Dict, Any, List


DEFAULT_CONFIG = {'conf': 0.25, 'iou': 0.45, 'agnostic_nms': False, 'max_det': 1000}
BOX_KEYS = ['xmin', 'ymin', 'xmax', 'ymax']

class EndpointHandler():
    def __init__(self, path=""):
        self.model = YOLO('ultralyticsplus/yolov8s')

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
       data args:
            image: image path to segment
            config: (conf - NMS confidence threshold,
                     iou - NMS IoU threshold,
                     agnostic_nms - NMS class-agnostic: True / False,
                     max_det - maximum number of detections per image)
      Return:
            A :obj: `dict` | `dict`: {scores, labels, boxes}
        """
        inputs = data.pop("inputs", data)
        input_config = inputs.pop("config", DEFAULT_CONFIG)
        config = {**DEFAULT_CONFIG, **input_config}
        
        if config is None:
            config = DEFAULT_CONFIG
        # Set model parameters
        self.model.overrides['conf'] = config.get('conf')
        self.model.overrides['iou'] = config.get('iou')
        self.model.overrides['agnostic_nms'] = config.get('agnostic_nms')
        self.model.overrides['max_det'] = config.get('max_det')   
        
        # Get label idx-to-name
        names = self.model.model.names
        
        # perform inference
        result = self.model.predict(inputs['image'])[0]
        prediction = []
        for score, label, box in zip(result.boxes.conf, result.boxes.cls, result.boxes.xyxy):
            item_score = score.item()
            item_label = names[int(label)]
            item_box = box.to(dtype=int).tolist()
            
            item_prediction = {
                'score': item_score,
                'label': item_label,
                'box': dict(zip(BOX_KEYS, item_box))
            }
            
            prediction.append(item_prediction)

        return prediction