import numpy as np
from .yolov5_postprocess import postprocess_
from base.preprocess import preprocess_, preprocess_hw
from base.postprocess import postprocess_hw
from base.inference import inference_

class Yolov5Runner:
    def __init__(self, model_path, model_id, input_shape, grid20_path, grid40_path, grid80_path, num_classes, conf_thres, iou_thres, top_k_num, anchors=None, label_mapping=None, vanish_point=0.0, filter_large_box_ratio=1.0, agnostic=0.0, **kwargs):
        # --------------- STEP 1 load model----------------
        from kneron_utils.model import load_model
        self.model, self._inference_type = load_model(model_path, **kwargs)

        # --------------- doing initialization-------------------        
        self.input_shape = input_shape
        # NOTE: stride can be an arg; currently hardcoded [8, 16, 32]
        max_s = 32
        assert input_shape[0] % max_s == 0 and input_shape[1] % max_s == 0
        self.model_id = str(model_id)
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.top_k_num = top_k_num
        self.num_classes = num_classes    
        self.vanish_point = vanish_point
        self.filter_large_box_ratio = filter_large_box_ratio
        self.agnostic = True if agnostic==1.0 else False
        self.grids = [np.load(grid80_path), np.load(grid40_path), np.load(grid20_path)]       
        self.anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] if anchors==None else anchors   

        # --------------- STEP 2 get all parameter as dict-----------------
        self.init_config = locals()
        self.init_config.update(kwargs)
        self.init_config['_type'] = self._inference_type                                                                                       
 
    def run(self, image, rectangle_list=None, **kwargs):
        
        # rectangle_list: List of rectangle [x1, y1, w, h, score]. where xy1=top-left.  
        if rectangle_list is not None:
            return self.run_rectangle_list(image, rectangle_list, **kwargs)

        # -----------------preprocess-------------------
        pre_config = {
            'input_shape': self.input_shape,
            'keep_ap': True,
            'upscale': True,
        }
        pre_config.update(self.init_config)
        pre_config.update(kwargs)
        img_data, pre_info = preprocess_hw(image, preprocess_, **pre_config)#img_data:rgb

        # -----------------do inference------------------
        infer_config = {'model': self.model}
        infer_config.update(self.init_config)
        infer_config.update(pre_config)
        outputs = inference_(img_data, **infer_config)
        
        # -----------------post process------------------
        post_config = {'model_id':self.model_id, 'conf_thres':self.conf_thres, 'iou_thres':self.iou_thres, 'top_k_num':self.top_k_num, 'grids':self.grids, 'num_classes':self.num_classes,  'vanish_point':self.vanish_point, 'filter_large_box_ratio':self.filter_large_box_ratio, 'agnostic':self.agnostic}
        post_config.update(self.init_config)
        post_config.update(pre_config)
        post_config.update(pre_info)
        post_config.update(kwargs)
        post_config['anchors'] = self.anchors
        
        outputs_onnx = postprocess_hw(outputs, postprocess_, **post_config)        
        # outputs_onnx : List of bboxes[x1, y1, w, h, score, class_id]. where xy1=top-left 
        label_mapping = post_config.get('label_mapping', None)
        if label_mapping is not None:
            for det in outputs_onnx:
                det[-1] = label_mapping[int(det[-1]-1)]
        return outputs_onnx
    
    def run_rectangle_list(self, image, rectangle_list, **kwargs):
        outputs_onnx_list = []
        # 'rectangle_list': List of 'rectangle' [x1, y1, w, h, score]. where xy1=top-left.
        # 'rectangle_list' = [[x1, y1, w, h, score], [x1, y1, w, h, score], ...]. 
        for ii, rectangle in enumerate(rectangle_list):
            #  -----------------preprocess------------------
            pre_config = {
                'input_shape': self.input_shape,
                'keep_ap': True,
                'upscale': True,
                'rectangle': rectangle,
            }
            pre_config.update(self.init_config)
            pre_config.update(kwargs)
            # 'img_data' is cropped from 'image' with 'rectangle' [x1, y1, w, h, score]
            img_data, pre_info = preprocess_hw(image, preprocess_, **pre_config)

            # -----------------do inference------------------
            infer_config = {'model': self.model}
            infer_config.update(self.init_config)
            infer_config.update(pre_config)
            outputs = inference_(img_data, **infer_config)

            # -----------------post process------------------
            post_config = {'model_id':self.model_id, 'conf_thres':self.conf_thres, 'iou_thres':self.iou_thres, 'top_k_num':self.top_k_num, 'grids':self.grids, 'num_classes':self.num_classes,  'vanish_point':self.vanish_point, 'filter_large_box_ratio':self.filter_large_box_ratio, 'agnostic':self.agnostic}
            post_config.update(self.init_config)
            post_config.update(pre_config)
            post_config.update(pre_info)
            post_config.update(kwargs)
            post_config['anchors'] = self.anchors
            outputs_onnx = postprocess_hw(outputs, postprocess_, **post_config) 
            # 'outputs_onnx' : List of bboxes[x1, y1, w, h, score, class_id]. where xy1=top-left 
            # 'outputs_onnx' = [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...]. 
            if len(outputs_onnx)==0: continue
            label_mapping = post_config.get('label_mapping', None)
            if label_mapping is not None:
                for det in outputs_onnx:
                    det[-1] = label_mapping[int(det[-1]-1)]
            outputs_onnx_list.append(outputs_onnx)
        # 'outputs_onnx_list' : List of outputs_onnx [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...].
        # 'outputs_onnx_list' = [ [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...],
        #                         [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...],
        #                         ...
        #                         [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...] ]. 
        return outputs_onnx_list