"""
This class implements the YoloDetector. It is used to detect elements in a
given image or batch of images. It uses YOLOv6 as the backbone of the detector.

Author: Eric Canas.
Github: https://github.com/Eric-Canas
Email: eric@ericcanas.com
Date: 09-07-2022
"""

import numpy as np

from vid2info.inference.detection.config import DETECTOR_MIN_CONFIDENCE_TH, IOU_NMS_TH, \
    MAX_DETECTIONS_PER_IMAGE, DETECTOR_IMAGE_SIZE, HALF_MODE, DATASET_YML, NORMALIZE_OUTPUT
from vid2info.inference.config import INFERENCE_DEVICE, XMIN, YMIN, YMAX, XMAX, CONFIDENCE, CLASS_ID, CLASS_NAME, BBOX_XYXY, \
    UINT8_TO_FLOAT
from vid2info.inference.utils import raw_yolo_bboxes_to_xyxy_co_cl
from yolov7.models.common import Conv
from yolov6.data.data_augment import letterbox
from yolov6.utils.nms import non_max_suppression

import os
import torch
import math
from warnings import warn

class __YoloDetector:
    def __init__(self, dataset_yml : str | None = DATASET_YML,
                 confidence_th : float = DETECTOR_MIN_CONFIDENCE_TH, nms_th : float = IOU_NMS_TH,
                 max_dets : int = MAX_DETECTIONS_PER_IMAGE, inference_device : str = INFERENCE_DEVICE,
                 image_size : int = DETECTOR_IMAGE_SIZE, half_mode = HALF_MODE,
                 merge_batches: bool = False, agnostic_nms: bool = False):

        """
        Initialize the YoloDetector (YoloV6).

        :param weights: String. Path to the weights file. It must be the .pt file generated by the YOLOv6 tool.
        :param dataset_yml: String. Path to the dataset .yaml file used for training the model in the YOLOv6 tool. It
                            must contain the 'nc' and 'names' fields.
        :param confidence_th: Float. The confidence threshold to use for assuming that a detection is valid. Take
                                into account that, when the detector is used within the pipeline together with the tracker,
                                this threshold will be used only for filtering high confidence detections with non-max
                                suppression, but it will also internally return everything below it. (To let the tracker
                                make its magic).
        :param nms_th: Float. The IoU threshold to use for non-max suppression.
        :param max_dets: Integer (greater than 0). The maximum number of detections to return by frame.
        :param inference_device: String ('cpu', 'cuda' or 'cuda:[int]'). The device to use for inference.
        :param image_size: Integer. The size of the images to use for inference.
        :param half_mode: Boolean. Whether to use half precision or not. Half precision only works with CUDA.
        :param merge_batches: Boolean. Whether to merge batches or not. If it is True, and the input is a batch
                                of images, predictions will be merged as if they came from a single image. It is only
                                recommended for very static scenes.
        :param agnostic_nms: Boolean. Whether to use agnostic nms or not. If it is True, the detector will use the
                                agnostic nms algorithm.
        """
        if inference_device.startswith('cuda'):
            assert torch.cuda.is_available(), 'CUDA is not available.'
        else:
            assert inference_device == 'cpu', f'Inference device must be cpu or cuda. Got {inference_device}.'
            assert half_mode is False, f'Half mode must be False for CPU inference.'

        if dataset_yml is not None:
            assert os.path.isfile(dataset_yml), f'Dataset yml file {dataset_yml} not found.'
            from yolov6.utils.events import load_yaml
            self.dataset_info = load_yaml(dataset_yml)
            self.class_names = self.dataset_info['names']
        else:
            self.dataset_info, self.class_names = None, None

        self.confidence_th = confidence_th
        self.nms_th = nms_th
        self.max_dets = max_dets
        self.image_size = image_size
        self.half_mode = half_mode
        self.inference_device = torch.device(inference_device)
        self.half = half_mode
        self.merge_batches = merge_batches
        self.agnostic_nms = agnostic_nms


        self.image_size = image_size
        # It must be set when inheriting this class.
        self.model = None
        self.stride = None

    def prepare_model(self):
        """
        Prepare the model for inference.
        """
        assert self.model is not None, "The model must be initialized before calling prepare_model."
        self.model = self.model

        for m in self.model.modules():
            if type(m) in (torch.nn.Hardswish, torch.nn.LeakyReLU, torch.nn.ReLU, torch.nn.ReLU6, torch.nn.SiLU):
                m.inplace = True  # pytorch 1.7.0 compatibility
            elif type(m) is torch.nn.Upsample:
                print('Upsample')
                m.recompute_scale_factor = None  # torch 1.11.0 compatibility
            elif type(m) is Conv:
                m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility

        input_sample = torch.zeros(1, len('RGB'), self.image_size, self.image_size, dtype=torch.float32,
                                   device=self.inference_device)
        self._predict(image=input_sample)
        if self.half:
            self.model = self.model.half()
        self.stride = self.model.stride
        if type(self.stride) is torch.Tensor:
            self.stride = self.stride[-1].item()

        self.image_size = self.check_img_size(self.image_size, s=self.stride)  # check image size

    def _predict(self, image : torch.Tensor) -> torch.Tensor:
        """
        Perform inference on an image.
        :param image: A tensor of shape (1, 3, h, w).
        :return: A tensor of shape (n_boxes, 5 + n_classes) containing the bounding boxes and the confidence for each
                    class.
        """
        assert self.model is not None, 'The model has not been initialized.'
        return self.model(image)

    def detect(self, image : np.ndarray, input_is_bgr : bool = False, classes_to_keep : None | list[int] | tuple[int] = None,
               return_as_dict : bool = False, return_normalized : bool = NORMALIZE_OUTPUT,
               return_all_below_th : float = 0.) -> list[dict] | np.ndarray:
        """
        Detect objects in a given image.
        NOTE: 99.99...% of the time, this function takes is spent in the inference. The unique way to
        make it faster is to batch images or, ideally, to use a GPU.
        :param image: Numpy.ndarray of shape (H, W, C). The image to detect objects in.
        :param input_is_bgr: Boolean. Whether the input image is in BGR format or not (RGB).
        :param classes_to_keep: List of integers. The classes to keep in the output.
        :return: Numpy.ndarray of shape (N, 6) or List of Dictionaries. The detected objects.
        """
        with torch.no_grad():
            image = self.preprocess_image(image=image, input_is_bgr=input_is_bgr)
            input_h, input_w = image.shape[-2:]
            if image.ndim == 3:
                image = image[None]
            assert image.ndim == 4, f'Image must be 4D. Got {image.shape}.'
            # Predictions here are in the format (x_center, y_center, width, height, confidence, class_id)
            pred = self._predict(image=image)
            if self.merge_batches:
                #pred = pred.reshape(1, -1, pred.shape[-1])
                best_conf_idx = torch.argmax(torch.sum(pred[..., 4:], axis=-1), axis=0)
                pred = pred[best_conf_idx, range(len(best_conf_idx))][None]
            below_th = []
            # Keeping classes below the confidence threshold could help the tracker.
            if return_all_below_th > 0.:
                below_th = pred[pred[..., 4] < return_all_below_th]
                below_th = raw_yolo_bboxes_to_xyxy_co_cl(preds=below_th)

            pred = non_max_suppression(prediction=pred, conf_thres=self.confidence_th, iou_thres=self.nms_th,
                                       classes=classes_to_keep, agnostic=self.agnostic_nms, multi_label=False, max_det=self.max_dets)
            if len(below_th) > 0:
                pred = (torch.cat((pred[0], below_th), dim=0),)
            assert len(pred) == 1, f'Expected 1 prediction. Got {len(pred)}.'
            pred = self.post_process_output_detection(prediction=pred[0], input_h=input_h, input_w=input_w,
                                                      as_dict=return_as_dict, normalize=return_normalized)
        return pred

    def post_process_output_detection(self, prediction : torch.Tensor, input_h : int, input_w : int,
                                      normalize : bool = True, as_dict : bool = False) -> list[dict] | np.ndarray:
        """
        Postprocess the output of the detector. Return a dictionary with the detected objects.

        :param prediction: List of tensors. The output of the detector.
        :param input_h: Integer. The height of the input image.
        :param input_w: Integer. The width of the input image.
        :param as_dict: Boolean. Whether to return the output as a dictionary or not.
        :return: List of Dictionaries or numpy.ndarray. The detected objects.
        As a dictionary, the keys are  ('xmin', 'ymin', 'xmax', 'ymax', 'class_id', 'class_name', 'confidence',
                                        'bbox_xyxy').
        As a numpy.ndarray, the shape is (N, 6). The columns are: (xmin, ymin, xmax, ymax, confidence, class_id).
        If normalize, all coordinates are in the range [0., 1.]
        """

        if normalize:
            # Convert from the bbox from (w1, h1, w2, h2) to the range [0., 1.] (w'1, h'1, w'2, h'2).
            prediction[:, [0, 2]] /= input_w
            prediction[:, [1, 3]] /= input_h
        if as_dict:
            prediction = prediction.tolist()
            prediction = [{XMIN: xmin, YMIN: ymin, XMAX: xmax, YMAX: ymax,
                           BBOX_XYXY: (xmin, ymin, xmax, ymax), CONFIDENCE: confidence,
                           CLASS_ID: int(class_id), CLASS_NAME: self.class_names[int(class_id)]}
                          for (xmin, ymin, xmax, ymax, confidence, class_id) in prediction]
        else:
            if self.inference_device.type != 'cpu':
                prediction = prediction.detach().cpu()
            prediction = prediction.numpy()

        return prediction

    def preprocess_image(self, image: np.ndarray, input_is_bgr = False) -> torch.Tensor:
        """
        Preprocess an image. It takes a raw RGB image and transforms it into the format that YOLOv6 expects.

        :param image: Numpy.ndarray of shape (H, W, 3). The image to preprocess. It must be in RGB or BGR format.
        :param input_is_bgr: Boolean. Whether the input image is in BGR format or not (RGB).

        :return: Tensor of shape (1, 3, H, W). The preprocessed image ready to be fed to the model.
        """
        assert image.dtype == np.uint8, f'Image must be uint8. Got {image.dtype}.'
        assert image.ndim in (3, 4), f'Image must be 3D. Got {image.shape}.'
        assert image.shape[-1] == 3, f'Image must have 3 channels (Format [B]HWC). Got {image.shape}.'

        if image.ndim == 3:
            image = letterbox(image, new_shape=self.image_size, stride=self.stride)[0]
            # HWC to CHW
            image = image.transpose((2, 0, 1))
        else:
            image = np.stack([letterbox(n_image, new_shape=self.image_size, stride=self.stride)[0] for n_image in image])
            # HWC to CHW
            image = image.transpose((0, 3, 1, 2))

        if input_is_bgr:
            # BGR to RGB
            image = image[..., ::-1, :, :]
        image = torch.tensor(np.ascontiguousarray(image), device=self.inference_device,
                             dtype=torch.float16 if self.half else torch.float32).to(self.inference_device)
        image *= UINT8_TO_FLOAT  # 0 - 255 to 0.0 - 1.0
        return image

    def check_img_size(self, image_size, s=32, floor=0):
        """
        Check if the image size is valid. If not, return the nearest valid size.

        :param image_size: Integer. The image size to check.
        :param s: Integer. The stride of the model
        :param floor: Integer. The minimum image size.

        :return: Integer. The valid image size.
        """
        def make_divisible(x, divisor):
            # Upward revision the value x to make it evenly divisible by the divisor.
            return math.ceil(x / divisor) * divisor
        """Make sure image size is a multiple of stride s in each dimension, and return a new shape list of image."""
        if isinstance(image_size, int):  # integer i.e. img_size=640
            new_size = max(make_divisible(image_size, int(s)), floor)
        elif isinstance(image_size, list) or isinstance(image_size, tuple):  # list i.e. img_size=[640, 480]
            new_size = [max(make_divisible(x, int(s)), floor) for x in image_size]
        else:
            raise Exception(f"Unsupported type of img_size: {type(image_size)}")

        if new_size != image_size:
            warn(f'image_size {image_size} must be multiple of max stride {s}, updating to {new_size}')
        return new_size if isinstance(image_size, list) else [new_size] * 2