import mxnet as mx
import numpy as np
import random
import cv2
from mxnet.executor_manager import _split_input_slice

from ..config import config
from ..io.image import tensor_vstack
from ..io.rpn import get_crop_batch, AA


class CropLoader(mx.io.DataIter):
    def __init__(self, feat_sym, roidb, batch_size=1, shuffle=False, ctx=None, work_load_list=None,
                 aspect_grouping=False):
        """
        This Iter will provide roi data to Fast R-CNN network
        :param feat_sym: to infer shape of assign_output
        :param roidb: must be preprocessed
        :param batch_size: must divide BATCH_SIZE(128)
        :param shuffle: bool
        :param ctx: list of contexts
        :param work_load_list: list of work load
        :param aspect_grouping: group images with similar aspects
        :return: AnchorLoader
        """
        super(CropLoader, self).__init__()

        # save parameters as properties
        self.feat_sym = feat_sym
        self.roidb = roidb
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.ctx = ctx
        if self.ctx is None:
            self.ctx = [mx.cpu()]
        self.work_load_list = work_load_list
        self.aspect_grouping = aspect_grouping
        self.feat_stride = config.RPN_FEAT_STRIDE

        # infer properties from roidb
        self.size = len(roidb)
        self.index = np.arange(self.size)

        self.data_name = ['data']
        self.label_name = []
        prefixes = ['face']
        if config.HEAD_BOX:
            prefixes.append('head')
        names = []
        for prefix in prefixes:
            names += [prefix + '_label', prefix + '_bbox_target', prefix + '_bbox_weight']
            if prefix == 'face' and config.FACE_LANDMARK:
                names += [prefix + '_landmark_target', prefix + '_landmark_weight']

        for stride in self.feat_stride:
            for n in names:
                k = "%s_stride%d" % (n, stride)
                self.label_name.append(k)
        if config.CASCADE > 0:
            self.label_name.append('gt_boxes')

        # status variable for synchronization between get_data and get_label
        self.cur = 0
        self.batch = None
        self.data = None
        self.label = None
        # infer shape
        feat_shape_list = []
        _data_shape = [('data', (1, 3, max([v[1] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]
        _data_shape = dict(_data_shape)
        for i in range(len(self.feat_stride)):
            _, feat_shape, _ = self.feat_sym[i].infer_shape(**_data_shape)
            feat_shape = [int(i) for i in feat_shape[0]]
            feat_shape_list.append(feat_shape)
        self.aa = AA(feat_shape_list)

        self._debug = False
        self._debug_id = 0
        self._times = [0.0, 0.0, 0.0, 0.0]

        # get first batch to fill in provide_data and provide_label
        self.reset()
        self.get_batch()

    @property
    def provide_data(self):
        return [(k, v.shape) for k, v in zip(self.data_name, self.data)]

    @property
    def provide_label(self):
        return [(k, v.shape) for k, v in zip(self.label_name, self.label)]

    def reset(self):
        self.cur = 0
        if self.shuffle:
            np.random.shuffle(self.index)

    def iter_next(self):
        return self.cur + self.batch_size <= self.size

    def next(self):
        if self.iter_next():
            self.get_batch()
            self.cur += self.batch_size
            return mx.io.DataBatch(data=self.data, label=self.label,
                                   pad=self.getpad(), index=self.getindex(),
                                   provide_data=self.provide_data, provide_label=self.provide_label)
        else:
            raise StopIteration

    def getindex(self):
        return self.cur / self.batch_size

    def getpad(self):
        if self.cur + self.batch_size > self.size:
            return self.cur + self.batch_size - self.size
        else:
            return 0

    def infer_shape(self, max_data_shape=None, max_label_shape=None):
        """ Return maximum data and label shape for single gpu """
        if max_data_shape is None:
            max_data_shape = []
        if max_label_shape is None:
            max_label_shape = []
        max_shapes = dict(max_data_shape + max_label_shape)
        input_batch_size = max_shapes['data'][0]
        dummy_boxes = np.zeros((0, 5))
        dummy_info = [[max_shapes['data'][2], max_shapes['data'][3], 1.0]]
        dummy_label = {'gt_boxes': dummy_boxes}
        dummy_blur = np.zeros((0,))
        dummy_label['gt_blur'] = dummy_blur

        label_dict = {}
        if config.HEAD_BOX:
            head_label_dict = self.aa.assign_anchor_fpn(dummy_label, dummy_info, False, prefix='head')
            label_dict.update(head_label_dict)

        if config.FACE_LANDMARK:
            dummy_landmarks = np.zeros((0, 5, 3))
            dummy_label['gt_landmarks'] = dummy_landmarks
        face_label_dict = self.aa.assign_anchor_fpn(dummy_label, dummy_info, config.FACE_LANDMARK, prefix='face')
        label_dict.update(face_label_dict)
        if config.CASCADE > 0:
            label_dict['gt_boxes'] = np.zeros((0, config.TRAIN.MAX_BBOX_PER_IMAGE, 5), dtype=np.float32)

        label_list = []
        for k in self.label_name:
            label_list.append(label_dict[k])
        label_shape = [(k, tuple([input_batch_size] + list(v.shape[1:]))) for k, v in zip(self.label_name, label_list)]
        return max_data_shape, label_shape

    def get_batch(self):
        # slice roidb
        cur_from = self.cur
        cur_to = min(cur_from + self.batch_size, self.size)
        assert cur_to == cur_from + self.batch_size
        roidb = [self.roidb[self.index[i]] for i in range(cur_from, cur_to)]

        # decide multi device slice
        work_load_list = self.work_load_list
        ctx = self.ctx
        if work_load_list is None:
            work_load_list = [1] * len(ctx)
        assert isinstance(work_load_list, list) and len(work_load_list) == len(ctx), \
            "Invalid settings for work load. "
        slices = _split_input_slice(self.batch_size, work_load_list)

        # get testing data for multigpu
        data_list = []
        label_list = []
        for islice in slices:
            iroidb = [roidb[i] for i in range(islice.start, islice.stop)]
            data, label = get_crop_batch(iroidb)
            data_list += data
            label_list += label

        select_stride = 0
        if config.RANDOM_FEAT_STRIDE:
            select_stride = random.choice(config.RPN_FEAT_STRIDE)

        for data, label in zip(data_list, label_list):
            data_shape = {k: v.shape for k, v in data.items()}
            del data_shape['im_info']
            feat_shape_list = []
            for s in range(len(self.feat_stride)):
                _, feat_shape, _ = self.feat_sym[s].infer_shape(**data_shape)
                feat_shape = [int(i) for i in feat_shape[0]]
                feat_shape_list.append(feat_shape)
            im_info = data['im_info']
            gt_boxes = label['gt_boxes']
            gt_label = {'gt_boxes': gt_boxes}
            if config.USE_BLUR:
                gt_blur = label['gt_blur']
                gt_label['gt_blur'] = gt_blur
            if self._debug:
                img = data['data'].copy()[0].transpose((1, 2, 0))[:, :, ::-1].copy()

                box = label['gt_boxes'].copy()[0][0:4].astype(np.int)
                cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
                filename = './debugout/%d.png' % (self._debug_id)
                cv2.imwrite(filename, img)
                self._debug_id += 1
            label_dict = {}
            if config.HEAD_BOX:
                head_label_dict = self.aa.assign_anchor_fpn(gt_label, im_info, False, prefix='head',
                                                            select_stride=select_stride)
                label_dict.update(head_label_dict)
            if config.FACE_LANDMARK:
                gt_landmarks = label['gt_landmarks']
                gt_label['gt_landmarks'] = gt_landmarks
            face_label_dict = self.aa.assign_anchor_fpn(gt_label, im_info, config.FACE_LANDMARK, prefix='face',
                                                        select_stride=select_stride)
            label_dict.update(face_label_dict)

            if config.CASCADE > 0:
                pad_gt_boxes = np.empty((1, config.TRAIN.MAX_BBOX_PER_IMAGE, 5), dtype=np.float32)
                pad_gt_boxes.fill(-1)
                pad_gt_boxes[0, 0:gt_boxes.shape[0], :] = gt_boxes
                label_dict['gt_boxes'] = pad_gt_boxes

            for k in self.label_name:
                label[k] = label_dict[k]

        all_data = dict()
        for key in self.data_name:
            all_data[key] = tensor_vstack([batch[key] for batch in data_list])

        all_label = dict()
        for key in self.label_name:
            pad = 0 if key.startswith('bbox_') else -1
            all_label[key] = tensor_vstack([batch[key] for batch in label_list], pad=pad)

        self.data = [mx.nd.array(all_data[key]) for key in self.data_name]
        self.label = [mx.nd.array(all_label[key]) for key in self.label_name]
