# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import numpy as np
import math

class RandSampler(object):
    """
    Random sampler base class, used for data augmentation

    Parameters:
    ----------
    max_trials : int
        maximum trials, if exceed this number, give up anyway
    max_sample : int
        maximum random crop samples to be generated
    """
    def __init__(self, max_trials, max_sample):
        assert max_trials > 0
        self.max_trials = int(max_trials)
        assert max_sample >= 0
        self.max_sample = int(max_sample)

    def sample(self, label):
        """
        Interface for calling sampling function

        Parameters:
        ----------
        label : numpy.array (n x 5 matrix)
            ground-truths

        Returns:
        ----------
        list of (crop_box, label) tuples, if failed, return empty list []
        """
        return NotImplementedError


class RandCropper(RandSampler):
    """
    Random cropping original images with various settings

    Parameters:
    ----------
    min_scale : float
        minimum crop scale, (0, 1]
    max_scale : float
        maximum crop scale, (0, 1], must larger than min_scale
    min_aspect_ratio : float
        minimum crop aspect ratio, (0, 1]
    max_aspect_ratio : float
        maximum crop aspect ratio, [1, inf)
    min_overlap : float
        hreshold of minimum overlap between a rand crop and any gt
    max_trials : int
        maximum trials, if exceed this number, give up anyway
    max_sample : int
        maximum random crop samples to be generated
    """
    def __init__(self, min_scale=1., max_scale=1.,
                 min_aspect_ratio=1., max_aspect_ratio=1.,
                 min_overlap=0., max_trials=50, max_sample=1):
        super(RandCropper, self).__init__(max_trials, max_sample)
        assert min_scale <= max_scale, "min_scale must <= max_scale"
        assert 0 < min_scale and min_scale <= 1, "min_scale must in (0, 1]"
        assert 0 < max_scale and max_scale <= 1, "max_scale must in (0, 1]"
        self.min_scale = min_scale
        self.max_scale = max_scale
        assert 0 < min_aspect_ratio and min_aspect_ratio <= 1, "min_ratio must in (0, 1]"
        assert 1 <= max_aspect_ratio , "max_ratio must >= 1"
        self.min_aspect_ratio = min_aspect_ratio
        self.max_aspect_ratio = max_aspect_ratio
        assert 0 <= min_overlap and min_overlap <= 1, "min_overlap must in [0,1]"
        self.min_overlap = min_overlap

        self.config = {'gt_constraint' : 'center'}

    def sample(self, label):
        """
        generate random cropping boxes according to parameters
        if satifactory crops generated, apply to ground-truth as well

        Parameters:
        ----------
        label : numpy.array (n x 5 matrix)
            ground-truths

        Returns:
        ----------
        list of (crop_box, label) tuples, if failed, return empty list []
        """
        samples = []
        count = 0
        for trial in range(self.max_trials):
            if count >= self.max_sample:
                return samples
            scale = np.random.uniform(self.min_scale, self.max_scale)
            min_ratio = max(self.min_aspect_ratio, scale * scale)
            max_ratio = min(self.max_aspect_ratio, 1. / scale / scale)
            ratio = math.sqrt(np.random.uniform(min_ratio, max_ratio))
            width = scale * ratio
            height = scale / ratio
            left = np.random.uniform(0., 1 - width)
            top = np.random.uniform(0., 1 - height)
            rand_box = (left, top, left + width, top + height)
            valid_mask = np.where(label[:, 0] > -1)[0]
            gt = label[valid_mask, :]
            ious = self._check_satisfy(rand_box, gt)
            if ious is not None:
                # transform gt labels after crop, discard bad ones
                l, t, r, b = rand_box
                new_gt_boxes = []
                new_width = r - l
                new_height = b - t
                for i in range(valid_mask.size):
                    if ious[i] > 0:
                        xmin = max(0., (gt[i, 1] - l) / new_width)
                        ymin = max(0., (gt[i, 2] - t) / new_height)
                        xmax = min(1., (gt[i, 3] - l) / new_width)
                        ymax = min(1., (gt[i, 4] - t) / new_height)
                        new_gt_boxes.append([gt[i, 0], xmin, ymin, xmax, ymax])
                if not new_gt_boxes:
                    continue
                new_gt_boxes = np.array(new_gt_boxes)
                label = np.lib.pad(new_gt_boxes,
                    ((0, label.shape[0]-new_gt_boxes.shape[0]), (0,0)), \
                    'constant', constant_values=(-1, -1))
                samples.append((rand_box, label))
                count += 1
        return samples

    def _check_satisfy(self, rand_box, gt_boxes):
        """
        check if overlap with any gt box is larger than threshold
        """
        l, t, r, b = rand_box
        num_gt = gt_boxes.shape[0]
        ls = np.ones(num_gt) * l
        ts = np.ones(num_gt) * t
        rs = np.ones(num_gt) * r
        bs = np.ones(num_gt) * b
        mask = np.where(ls < gt_boxes[:, 1])[0]
        ls[mask] = gt_boxes[mask, 1]
        mask = np.where(ts < gt_boxes[:, 2])[0]
        ts[mask] = gt_boxes[mask, 2]
        mask = np.where(rs > gt_boxes[:, 3])[0]
        rs[mask] = gt_boxes[mask, 3]
        mask = np.where(bs > gt_boxes[:, 4])[0]
        bs[mask] = gt_boxes[mask, 4]
        w = rs - ls
        w[w < 0] = 0
        h = bs - ts
        h[h < 0] = 0
        inter_area = h * w
        union_area = np.ones(num_gt) * max(0, r - l) * max(0, b - t)
        union_area += (gt_boxes[:, 3] - gt_boxes[:, 1]) * (gt_boxes[:, 4] - gt_boxes[:, 2])
        union_area -= inter_area
        ious = inter_area / union_area
        ious[union_area <= 0] = 0
        max_iou = np.amax(ious)
        if max_iou < self.min_overlap:
            return None
        # check ground-truth constraint
        if self.config['gt_constraint'] == 'center':
            for i in range(ious.shape[0]):
                if ious[i] > 0:
                    gt_x = (gt_boxes[i, 1] + gt_boxes[i, 3]) / 2.0
                    gt_y = (gt_boxes[i, 2] + gt_boxes[i, 4]) / 2.0
                    if gt_x < l or gt_x > r or gt_y < t or gt_y > b:
                        return None
        elif self.config['gt_constraint'] == 'corner':
            for i in range(ious.shape[0]):
                if ious[i] > 0:
                    if gt_boxes[i, 1] < l or gt_boxes[i, 3] > r \
                        or gt_boxes[i, 2] < t or gt_boxes[i, 4] > b:
                        return None
        return ious


class RandPadder(RandSampler):
    """
    Random cropping original images with various settings

    Parameters:
    ----------
    min_scale : float
        minimum crop scale, [1, inf)
    max_scale : float
        maximum crop scale, [1, inf), must larger than min_scale
    min_aspect_ratio : float
        minimum crop aspect ratio, (0, 1]
    max_aspect_ratio : float
        maximum crop aspect ratio, [1, inf)
    min_gt_scale : float
        minimum ground-truth scale to be satisfied after padding,
        either width or height, [0, 1]
    max_trials : int
        maximum trials, if exceed this number, give up anyway
    max_sample : int
        maximum random crop samples to be generated
    """
    def __init__(self, min_scale=1., max_scale=1., min_aspect_ratio=1., \
                 max_aspect_ratio=1., min_gt_scale=.01, max_trials=50,
                 max_sample=1):
        super(RandPadder, self).__init__(max_trials, max_sample)
        assert min_scale <= max_scale, "min_scale must <= max_scale"
        assert min_scale >= 1, "min_scale must in (0, 1]"
        self.min_scale = min_scale
        self.max_scale = max_scale
        assert 0 < min_aspect_ratio and min_aspect_ratio <= 1, "min_ratio must in (0, 1]"
        assert 1 <= max_aspect_ratio , "max_ratio must >= 1"
        self.min_aspect_ratio = min_aspect_ratio
        self.max_aspect_ratio = max_aspect_ratio
        assert 0 <= min_gt_scale and min_gt_scale <= 1, "min_gt_scale must in [0, 1]"
        self.min_gt_scale = min_gt_scale

    def sample(self, label):
        """
        generate random padding boxes according to parameters
        if satifactory padding generated, apply to ground-truth as well

        Parameters:
        ----------
        label : numpy.array (n x 5 matrix)
            ground-truths

        Returns:
        ----------
        list of (crop_box, label) tuples, if failed, return empty list []
        """
        samples = []
        count = 0
        for trial in range(self.max_trials):
            if count >= self.max_sample:
                return samples
            scale = np.random.uniform(self.min_scale, self.max_scale)
            min_ratio = max(self.min_aspect_ratio, scale * scale)
            max_ratio = min(self.max_aspect_ratio, 1. / scale / scale)
            ratio = math.sqrt(np.random.uniform(min_ratio, max_ratio))
            width = scale * ratio
            if width < 1:
                continue
            height = scale / ratio
            if height < 1:
                continue
            left = np.random.uniform(0., 1 - width)
            top = np.random.uniform(0., 1 - height)
            right = left + width
            bot = top + height
            rand_box = (left, top, right, bot)
            valid_mask = np.where(label[:, 0] > -1)[0]
            gt = label[valid_mask, :]
            new_gt_boxes = []
            for i in range(gt.shape[0]):
                xmin = (gt[i, 1] - left) / width
                ymin = (gt[i, 2] - top) / height
                xmax = (gt[i, 3] - left) / width
                ymax = (gt[i, 4] - top) / height
                new_size = min(xmax - xmin, ymax - ymin)
                if new_size < self.min_gt_scale:
                    new_gt_boxes = []
                    break
                new_gt_boxes.append([gt[i, 0], xmin, ymin, xmax, ymax])
            if not new_gt_boxes:
                continue
            new_gt_boxes = np.array(new_gt_boxes)
            label = np.lib.pad(new_gt_boxes,
                ((0, label.shape[0]-new_gt_boxes.shape[0]), (0,0)), \
                'constant', constant_values=(-1, -1))
            samples.append((rand_box, label))
            count += 1
        return samples
