"""
Object detection inference utilities
"""
import sys
import cv2
import numpy as np
import tensorflow as tf
from time import time

from .yolov5 import YOLOv5
from .efficientdet import EfficientDet
from .config import CFG


def preprocess(img, img_shape):
    """Preprocess frame for inference

    Args:
        img (np.ndarray): BGR image
        img_shape (tuple): target image size in first two dimensions (K, L)

    Returns:
        tf.Tensor: target image for inference
    """
    # resize image to target size
    img = cv2.resize(img, img_shape)
    # transform from <0, 255> range to <0, 1> and convert to tensorflow Tensor
    img = img[np.newaxis, ...].astype(np.float32) / 255
    img = tf.convert_to_tensor(img, dtype=tf.float32)
    return img


def postprocess(prediction, img_size):
    """Postprocess detection outputs

    Args:
        prediction (tuple): detector's outputs - tuple of np.ndarrays - (boxes, scores, classes)
        img_size (tuple): input image size in first two dimensions (K, L)

    Returns:
        np.ndarray, np.array, np.array: predicted bounding boxes and their scores and classes
    """
    boxes = prediction[0][0, :].numpy()
    scores = prediction[1][0, :].numpy()
    classes = prediction[2][0, :].numpy()

    # filter out zero score predictions
    mask = scores != 0
    boxes = boxes[mask, :]
    scores = scores[mask]
    classes = classes[mask].astype(np.int32)
    # transform normed boxes from <0, 1> ranges to image plane space
    boxes[:, :2] = boxes[:, :2] * img_size
    boxes[:, 2:] = boxes[:, 2:] * img_size
    return boxes, scores, classes


def draw_boxes(img, outs):
    """Draw detections on image

    Args:
        img (np.ndarray): BGR image to draw on 
        outs (tuple): detector's outputs - tuple of np.ndarrays - (boxes, scores, classes)
    """
    for box, score, cls in zip(*outs):
        x, y, w, h = box
        xy1 = (int(x-w/2), int(y-h/2))
        xy2 = (int(x+w/2), int(y+h/2))
        color = (191, 111, 39) if cls == 0 else (71, 58, 240)
        cv2.rectangle(img, xy1, xy2, color, 2)
        cv2.putText(img, f'{score:.2f}', (xy1[0], xy1[1] - 10), cv2.FONT_HERSHEY_DUPLEX, 1, color, 2)


class Timer:
    """Inference timer
    """
    def __init__(self, k=200):
        """
        Args:
            k (int, optional): number of samples to keep track of. Defaults to 200.
        """
        self.k = k
        self.vals = k*[0]
        self.idx = 0
        self.n = 0

    @staticmethod
    def time_method(method, *args, **kwargs):
        """Time latency of a method

        Args:
            method (callable): method to measure

        Returns:
            any, float: results of measure method and measure latency
        """
        t0 = time()
        result = method(*args, **kwargs)
        t1 = time()
        residual = t1 - t0
        return result, residual

    def update(self, value):
        """Update timer with last measured latency

        Args:
            value (float): measured latency
        """
        self.vals[self.idx] = value
        self.idx = (self.idx + 1) % self.k
        if self.vals[-1] == 0:
            self.n += 1

    def __repr__(self):
        return f'{1000*sum(self.vals)/self.n:.2f}ms'


class ExportModel:
    """Exported model with computational graph
    """
    def __init__(self, model):
        self.model = model
    
    @tf.function
    def f(self, x):
        return self.model(x)


def detect(model, image):
    """Model prediction with pre and post processing

    Args:
        model (ExportModel): exported object detector
        image (tf.Tensor): input image

    Returns:
        np.ndarray, np.array, np.array: postprocessed detector's outputs - boxes, scores, classes
    """
    preprocessed = preprocess(image, CFG.COMBINED_PARAMS.IMG_SHAPE)
    out = model.f(preprocessed)
    boxes, scores, classes = postprocess(out, image.shape[1::-1])
    return boxes, scores, classes


def timed_predict(model, tensor, timer):
    """Time model prediction

    Args:
        model (ExportModel): exported object detector
        tensor (tf.Tensor): preprocessed input tensor
        timer (Timer): timer object

    Returns:
        np.ndarray, np.array, np.array: raw detector's outputs - boxes, scores, classes
    """
    result, t = Timer.time_method(model.f, tensor)
    timer.update(t)
    sys.stdout.write(f'\rModel prediction: {timer}' + 10*' ')
    sys.stdout.flush()
    return result


def init_detector(model_cfg):
    """Load object detector for inference

    Args:
        model_cfg (str): object detector configuration, either 'yolov5-s', 'yolov5-m', 'yolov5-l'
                         or 'efficientdet-d0' up to 'efficientdet-d6', based on available trained weights

    Returns:
        ExportModel: Exported model for inference
    """
    img_shape = (CFG.COMBINED_PARAMS.IMG_SHAPE[0], CFG.COMBINED_PARAMS.IMG_SHAPE[1], 3)
    model_f = YOLOv5 if 'yolov5' in model_cfg else EfficientDet
    anchors = CFG.COMBINED_PARAMS.YOLO_ANCHORS if 'yolov5' in model_cfg else CFG.COMBINED_PARAMS.EFFDET_ANCHORS
    # construct model
    model = model_f(img_shape, 2, anchors, model=model_cfg, training=False)
    # build graph for one input image
    model.build((1, *img_shape))
    # load trained weights
    model.load_weights(CFG.COMBINED_WEIGHTS_PATH[model_cfg]).expect_partial()
    # export model for inference
    model = ExportModel(model)
    model.f(np.random.rand(1, *img_shape))
    return model
