import os
import khandy
import numpy as np
from .base import OnnxModel
from .base import check_image_dtype_and_shape


class InsectDetector(OnnxModel):
    def __init__(self):
        current_dir = os.path.dirname(os.path.abspath(__file__))
        model_path = os.path.join(
            current_dir,
            "__pycache__/MuGemSt/insecta/quarrying_insect_detector.onnx",
        )
        self.input_width = 640
        self.input_height = 640
        super(InsectDetector, self).__init__(model_path)

    def _preprocess(self, image):
        check_image_dtype_and_shape(image)
        # image size normalization
        image, scale, pad_left, pad_top = khandy.letterbox_image(
            image, self.input_width, self.input_height, 0, return_scale=True
        )
        # image channel normalization
        image = khandy.normalize_image_channel(image, swap_rb=True)
        # image dtype normalization
        image = khandy.rescale_image(image, "auto", np.float32)
        # to tensor
        image = np.transpose(image, (2, 0, 1))
        image = np.expand_dims(image, axis=0)
        return image, scale, pad_left, pad_top

    def _post_process(
        self, outputs_list, scale, pad_left, pad_top, conf_thresh, iou_thresh
    ):
        pred = outputs_list[0][0]
        pass_t = pred[:, 4] > conf_thresh
        pred = pred[pass_t]
        boxes = khandy.convert_boxes_format(pred[:, :4], "cxcywh", "xyxy")
        boxes = khandy.unletterbox_2d_points(boxes, scale, pad_left, pad_top, False)
        confs = np.max(pred[:, 5:] * pred[:, 4:5], axis=-1)
        classes = np.argmax(pred[:, 5:] * pred[:, 4:5], axis=-1)
        keep = khandy.non_max_suppression(boxes, confs, iou_thresh)
        return boxes[keep], confs[keep], classes[keep]

    def detect(self, image, conf_thresh=0.5, iou_thresh=0.5):
        image, scale, pad_left, pad_top = self._preprocess(image)
        outputs_list = self.forward(image)
        boxes, confs, classes = self._post_process(
            outputs_list,
            scale=scale,
            pad_left=pad_left,
            pad_top=pad_top,
            conf_thresh=conf_thresh,
            iou_thresh=iou_thresh,
        )
        return boxes, confs, classes