#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.
import copy
import os

import cv2
import random
import numpy as np
from pycocotools.coco import COCO

from yolox.data.dataloading import get_yolox_datadir
from yolox.data.datasets.datasets_wrapper import CacheDataset, cache_read_img
from yolox.data.datasets.coco import remove_useless_info
from yolox.data.data_augment import TrainTransform, augment_hsv, _mirror, ValTransform
from yolox.utils import xyxy2cxcywh


class COCODataset_custom(CacheDataset):
    """
    COCO dataset class.
    """

    def __init__(
        self,
        catids=None,
        data_dir=None,
        json_file="instances_train2017.json",
        name="train2017",
        img_size=(416, 416),
        preproc=None,
        cache=False,
        cache_type="ram",
        max_imgs=None
    ):
        """
        COCO dataset initialization. Annotation data are read into memory by COCO API.
        Args:
            data_dir (str): dataset root directory
            json_file (str): COCO json file name
            name (str): COCO data name (e.g. 'train2017' or 'val2017')
            img_size (int): target image size after pre-processing
            preproc: data augmentation strategy
        """
        if data_dir is None:
            data_dir = os.path.join(get_yolox_datadir(), "COCO")
        self.data_dir = data_dir
        self.json_file = json_file
        self.max_imgs = max_imgs
        self.catids = catids
        self.coco = COCO(os.path.join(self.data_dir, "annotations", self.json_file))
        ids = []
        for catid in catids:
            ids.extend(self.coco.getImgIds(catIds=[catid]))
        ids = set(ids)
        self.ids = list(ids)
        if self.max_imgs is not None:
            self.ids = self.ids[:self.max_imgs]
        print("num images: %s"%(len(self.ids)))
        self.num_imgs = len(self.ids)
        self.class_ids = sorted(self.coco.getCatIds(catIds=catids))
        self.cats = self.coco.loadCats(self.coco.getCatIds(catIds=catids))
        self._classes = tuple([c["name"] for c in self.cats])
        self.name = name
        self.img_size = img_size
        self.preproc = preproc
        self.annotations = self._load_coco_annotations()

        path_filename = [os.path.join(name, anno[3]) for anno in self.annotations]

        super().__init__(
            input_dimension=img_size,
            num_imgs=self.num_imgs,
            data_dir=data_dir,
            cache_dir_name=f"cache_{name}",
            path_filename=path_filename,
            cache=cache,
            cache_type=cache_type
        )

    def __len__(self):
        return self.num_imgs

    def _load_coco_annotations(self):
        return [self.load_anno_from_ids(_ids) for _ids in self.ids]

    def load_anno_from_ids(self, id_):
        im_ann = self.coco.loadImgs(id_)[0]
        width = im_ann["width"]
        height = im_ann["height"]
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=False)
        annotations = self.coco.loadAnns(anno_ids)
        objs = []
        for obj in annotations:
            x1 = np.max((0, obj["bbox"][0]))
            y1 = np.max((0, obj["bbox"][1]))
            x2 = np.min((width, x1 + np.max((0, obj["bbox"][2]))))
            y2 = np.min((height, y1 + np.max((0, obj["bbox"][3]))))
            if obj["area"] > 0 and x2 >= x1 and y2 >= y1 and obj["category_id"] in self.class_ids:
                obj["clean_bbox"] = [x1, y1, x2, y2]
                objs.append(obj)

        num_objs = len(objs)

        res = np.zeros((num_objs, 5))
        for ix, obj in enumerate(objs):
            cls = self.class_ids.index(obj["category_id"])
            res[ix, 0:4] = obj["clean_bbox"]
            res[ix, 4] = cls

        r = min(self.img_size[0] / height, self.img_size[1] / width)
        res[:, :4] *= r

        img_info = (height, width)
        resized_info = (int(height * r), int(width * r))

        file_name = (
            im_ann["file_name"]
            if "file_name" in im_ann
            else "{:012}".format(id_) + ".jpg"
        )

        return (res, img_info, resized_info, file_name)

    def load_anno(self, index):
        return self.annotations[index][0]

    def load_resized_img(self, index):
        img = self.load_image(index)
        r = min(self.img_size[0] / img.shape[0], self.img_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * r), int(img.shape[0] * r)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        return resized_img

    def load_image(self, index):
        file_name = self.annotations[index][3]

        img_file = os.path.join(self.data_dir, self.name, file_name)

        img = cv2.imread(img_file)
        assert img is not None, f"file named {img_file} not found"

        return img

    @cache_read_img(use_cache=True)
    def read_img(self, index):
        return self.load_resized_img(index)

    def pull_item(self, index):
        id_ = self.ids[index]
        label, origin_image_size, _, _ = self.annotations[index]
        img = self.read_img(index)

        return img, copy.deepcopy(label), origin_image_size, np.array([id_])

    @CacheDataset.mosaic_getitem
    def __getitem__(self, index):
        """
        One image / label pair for the given index is picked up and pre-processed.

        Args:
            index (int): data index

        Returns:
            img (numpy.ndarray): pre-processed image
            padded_labels (torch.Tensor): pre-processed label data.
                The shape is :math:`[max_labels, 5]`.
                each label consists of [class, xc, yc, w, h]:
                    class (float): class index.
                    xc, yc (float) : center of bbox whose values range from 0 to 1.
                    w, h (float) : size of bbox whose values range from 0 to 1.
            info_img : tuple of h, w.
                h, w (int): original shape of the image
            img_id (int): same as the input index. Used for evaluation.
        """
        img, target, img_info, img_id = self.pull_item(index)

        if self.preproc is not None:
            img, target = self.preproc(img, target, self.input_dim)
        return img, target, img_info, img_id


def ToNv12(image, code):
    I420 = cv2.cvtColor(image, code)
    height, width = image.shape[:2]

    U = I420[height:height+height//4, :]
    U = U.reshape((1, height//4*width))
    V = I420[height+height//4:height+height//2, :]
    V = V.reshape((1, height//4*width))
    UV = np.zeros((1, height//4*width*2))
    UV[:, 0::2] = U
    UV[:, 1::2] = V
    UV = UV.reshape((height//2, width))
    NV12 = np.zeros((height+height//2, width))
    NV12[0:height, :] = I420[0:height, :]
    NV12[height:, :] = UV
    return NV12

def nv12toyuv444(nv12_data, height, width):
    nv12_data = nv12_data.reshape((height+height//2, width))
    y = nv12_data[0:height, :]
    u = nv12_data[height:height+height//2, 0::2]
    u = u.repeat(2, axis=0).repeat(2, axis=1)
    v = nv12_data[height:height+height//2, 1::2]
    v = v.repeat(2, axis=0).repeat(2, axis=1)
    yuv444 = np.zeros((height, width, 3))
    yuv444[:, :, 0] = y
    yuv444[:, :, 1] = u
    yuv444[:, :, 2] = v
    return yuv444


def preproc_yuv444(img, input_size, swap=(2, 0, 1)):
    if len(img.shape) == 3:
        padded_img = np.zeros((input_size[0], input_size[1], 3), dtype=np.uint8)
    else:
        padded_img = np.zeros(input_size, dtype=np.uint8)

    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
    resized_img = cv2.resize(
        img,
        (int(img.shape[1] * r), int(img.shape[0] * r)),
        interpolation=cv2.INTER_LINEAR,
    ).astype(np.uint8)
    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
    image_height, image_width, _ = padded_img.shape
    padded_img = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB)
    padded_img = ToNv12(padded_img, cv2.COLOR_RGB2YUV_I420).astype(np.uint8)
    padded_img = nv12toyuv444(padded_img, image_height, image_width).astype(np.uint8)
    # nomalize
    padded_img = padded_img.astype(np.float32) / 255.0
    padded_img = padded_img.transpose(swap)
    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
    return padded_img, r


def preproc_rgb(img, input_size, swap=(2, 0, 1)):
    if len(img.shape) == 3:
        padded_img = np.zeros((input_size[0], input_size[1], 3), dtype=np.uint8)
    else:
        padded_img = np.zeros(input_size, dtype=np.uint8)

    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
    resized_img = cv2.resize(
        img,
        (int(img.shape[1] * r), int(img.shape[0] * r)),
        interpolation=cv2.INTER_LINEAR,
    ).astype(np.uint8)
    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
    image_height, image_width, _ = padded_img.shape
    padded_img = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB)
    # nomalize
    padded_img = padded_img.astype(np.float32) / 255.0
    padded_img = padded_img.transpose(swap)
    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
    return padded_img, r


class TrainTransform_yuv444(TrainTransform):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def __call__(self, image, targets, input_dim):
        boxes = targets[:, :4].copy()
        labels = targets[:, 4].copy()
        if len(boxes) == 0:
            targets = np.zeros((self.max_labels, 5), dtype=np.float32)
            image, r_o = preproc_yuv444(image, input_dim)
            return image, targets

        image_o = image.copy()
        targets_o = targets.copy()
        height_o, width_o, _ = image_o.shape
        boxes_o = targets_o[:, :4]
        labels_o = targets_o[:, 4]
        # bbox_o: [xyxy] to [c_x,c_y,w,h]
        boxes_o = xyxy2cxcywh(boxes_o)

        if random.random() < self.hsv_prob:
            augment_hsv(image)
        image_t, boxes = _mirror(image, boxes, self.flip_prob)
        height, width, _ = image_t.shape
        image_t, r_ = preproc_yuv444(image_t, input_dim)
        # boxes [xyxy] 2 [cx,cy,w,h]
        boxes = xyxy2cxcywh(boxes)
        boxes *= r_

        mask_b = np.minimum(boxes[:, 2], boxes[:, 3]) > 1
        boxes_t = boxes[mask_b]
        labels_t = labels[mask_b]

        if len(boxes_t) == 0:
            image_t, r_o = preproc_yuv444(image_o, input_dim)
            boxes_o *= r_o
            boxes_t = boxes_o
            labels_t = labels_o

        labels_t = np.expand_dims(labels_t, 1)

        targets_t = np.hstack((labels_t, boxes_t))
        padded_labels = np.zeros((self.max_labels, 5))
        padded_labels[range(len(targets_t))[: self.max_labels]] = targets_t[
            : self.max_labels
        ]
        padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)
        return image_t, padded_labels
    

class ValTransform_yuv444(ValTransform):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    # assume input is cv2 img for now
    def __call__(self, img, res, input_size):
        img, _ = preproc_yuv444(img, input_size, self.swap)
        return img, np.zeros((1, 5))
    

class TrainTransform_rgb(TrainTransform):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def __call__(self, image, targets, input_dim):
        boxes = targets[:, :4].copy()
        labels = targets[:, 4].copy()
        if len(boxes) == 0:
            targets = np.zeros((self.max_labels, 5), dtype=np.float32)
            image, r_o = preproc_rgb(image, input_dim)
            return image, targets

        image_o = image.copy()
        targets_o = targets.copy()
        height_o, width_o, _ = image_o.shape
        boxes_o = targets_o[:, :4]
        labels_o = targets_o[:, 4]
        # bbox_o: [xyxy] to [c_x,c_y,w,h]
        boxes_o = xyxy2cxcywh(boxes_o)

        if random.random() < self.hsv_prob:
            augment_hsv(image)
        image_t, boxes = _mirror(image, boxes, self.flip_prob)
        height, width, _ = image_t.shape
        image_t, r_ = preproc_rgb(image_t, input_dim)
        # boxes [xyxy] 2 [cx,cy,w,h]
        boxes = xyxy2cxcywh(boxes)
        boxes *= r_

        mask_b = np.minimum(boxes[:, 2], boxes[:, 3]) > 1
        boxes_t = boxes[mask_b]
        labels_t = labels[mask_b]

        if len(boxes_t) == 0:
            image_t, r_o = preproc_rgb(image_o, input_dim)
            boxes_o *= r_o
            boxes_t = boxes_o
            labels_t = labels_o

        labels_t = np.expand_dims(labels_t, 1)

        targets_t = np.hstack((labels_t, boxes_t))
        padded_labels = np.zeros((self.max_labels, 5))
        padded_labels[range(len(targets_t))[: self.max_labels]] = targets_t[
            : self.max_labels
        ]
        padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)
        return image_t, padded_labels
    

class ValTransform_rgb(ValTransform):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    # assume input is cv2 img for now
    def __call__(self, img, res, input_size):
        img, _ = preproc_rgb(img, input_size, self.swap)
        return img, np.zeros((1, 5))