from engines.inference_engine cimport InferenceEngine
from annotation cimport Detection
cimport constants_inf
import numpy as np
from PIL import Image
import cv2
import io
import os
import tempfile
import zipfile


cdef class CoreMLEngine(InferenceEngine):

    def __init__(self, model_bytes: bytes, max_batch_size: int = 1, **kwargs):
        InferenceEngine.__init__(self, model_bytes, max_batch_size, engine_name="coreml")
        import coremltools as ct

        model_path = kwargs.get('model_path')
        if model_path is None:
            model_path = self._extract_from_zip(model_bytes)

        self.model = ct.models.MLModel(
            model_path, compute_units=ct.ComputeUnit.ALL)
        spec = self.model.get_spec()

        img_input = spec.description.input[0]
        self.img_width = int(img_input.type.imageType.width)
        self.img_height = int(img_input.type.imageType.height)

        constants_inf.log(<str>f'CoreML model: {self.img_width}x{self.img_height}')

    @staticmethod
    def _extract_from_zip(model_bytes):
        tmpdir = tempfile.mkdtemp()
        buf = io.BytesIO(model_bytes)
        with zipfile.ZipFile(buf, 'r') as zf:
            zf.extractall(tmpdir)
        for item in os.listdir(tmpdir):
            if item.endswith('.mlpackage') or item.endswith('.mlmodel'):
                return os.path.join(tmpdir, item)
        raise ValueError("No .mlpackage or .mlmodel found in zip")

    cdef tuple get_input_shape(self):
        return <tuple>(self.img_height, self.img_width)

    cdef preprocess(self, list frames):
        frame = frames[0]
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb, (self.img_width, self.img_height))
        return Image.fromarray(resized)

    cdef run(self, input_data):
        predict = getattr(self.model, 'predict')
        return predict({
            'image': input_data,
            'iouThreshold': 0.45,
            'confidenceThreshold': 0.25,
        })

    cdef list postprocess(self, output, object ai_config):
        cdef int w = self.img_width
        cdef int h = self.img_height

        coords = output.get('coordinates', np.empty((0, 4), dtype=np.float32))
        confs = output.get('confidence', np.empty((0, 80), dtype=np.float32))

        cdef list[Detection] detections = []
        if coords.size == 0:
            return [detections]

        cx, cy, bw, bh = coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]
        x1 = (cx - bw / 2) * w
        y1 = (cy - bh / 2) * h
        x2 = (cx + bw / 2) * w
        y2 = (cy + bh / 2) * h

        class_ids = np.argmax(confs, axis=1)
        conf_values = np.max(confs, axis=1)

        for i in range(len(conf_values)):
            conf = round(float(conf_values[i]), 2)
            if conf < ai_config.probability_threshold:
                continue
            det_x1 = float(x1[i]) / w
            det_y1 = float(y1[i]) / h
            det_x2 = float(x2[i]) / w
            det_y2 = float(y2[i]) / h
            det_cx = (det_x1 + det_x2) / 2
            det_cy = (det_y1 + det_y2) / 2
            det_w = det_x2 - det_x1
            det_h = det_y2 - det_y1
            detections.append(Detection(det_cx, det_cy, det_w, det_h, int(class_ids[i]), conf))

        filtered = self.remove_overlapping(detections, ai_config.tracking_intersection_threshold)
        return [filtered]
