<gh_stars>0
from enum import Enum

import torch

from carbrain.io.exceptions import UnknownBoundingBoxEncodingException


class BoundingBoxEncoding(Enum):
    """Encoding for a BoundingBox object

    """
    ULLR = 1
    CCWH = 2
    NORMALIZED_ULLR = 3


class BoundingBoxes(object):
    """Class for storing bounding boxes

    The class can be used to store bounding boxes for, e.g., single video
    frames.

    Args:
        coords (numpy.ndarray or torch.tensor): The coordinates of the
            bounding boxes.
        classes (numpy.ndarray or torch.tensor, optional): Class indices for
            the supplied bounding boxes.
        encoding (BoundingBoxEncoding, optional): The encoding of the bounding
            boxes in the coords argument.
        encoding_params (numpy.ndarray or torch.tensors, optional): Parameter
            which accompanies the BoundingBoxEncoding encoding.
        **attributes (optional): keyword args for additional attributes. Each
            attribute must be a numpy.ndarray or a torch.tensor with the shape
            of the first dimension matching the number of bounding boxes,
            i.e., there is an attribute for each bounding box.
    """
    def __init__(self, coords, classes=None, encoding=BoundingBoxEncoding.ULLR, encoding_params=None, **attributes):
        self._coords = torch.tensor(coords, dtype=torch.float32)
        if classes is None:
            self._classes = None
        else:
            self._classes = torch.tensor(classes, dtype=torch.int32)
            if self._classes.shape[0] != self._coords.shape[0]:
                raise ValueError(
                    'You must either supply None as class attribute or an '
                    'array of as many elements as there are bounding box '
                    'coordinates entries.'
                )
        self._encoding = encoding
        if encoding_params is not None:
            self._encoding_params = torch.tensor(encoding_params, dtype=torch.float32)
        else:
            self._encoding_params = None

        self._attributes = []
        for attribute_name, attribute_value in attributes.items():
            attribute_value = torch.tensor(attribute_value)
            # make sure there is an attribute for each bounding box
            if attribute_value.shape[0] != self._coords.shape[0]:
                raise ValueError(
                    'The number of values in {} does not match the number '
                    'of bounding boxes'
                    .format(attribute_name)
                )
            self.__setattr__(attribute_name, attribute_value)
            self._attributes.append(attribute_name)

    @property
    def coords(self):
        return self._coords

    @property
    def classes(self):
        return self._classes

    @property
    def encoding(self):
        return self._encoding

    @property
    def encoding_params(self):
        return self._encoding_params

    def encode(self, encoding, encoding_params=None):
        """Encode the current bounding box coordinates

        Args:
            encoding (BoundingBoxEncoding): The encoding the bounding box
                coordinates shall be converted to.
            encoding_params (numpy.ndarray or torch.tensors, optional): Parameter
                which accompanies the BoundingBoxEncoding encoding.
        """
        if encoding != self._encoding:
            self.decode()
            if encoding == BoundingBoxEncoding.ULLR:
                # bounding box encoding is already in ULLR after decode
                pass
            elif encoding == BoundingBoxEncoding.CCWH:
                width_height = self._coords[:, 2:] - self._coords[:, :2]
                center = self._coords[:, :2] + width_height/2
                self._coords = torch.cat((center, width_height), dim=1)
                self._encoding = BoundingBoxEncoding.CCWH
                self._encoding_params = None
            elif encoding == BoundingBoxEncoding.NORMALIZED_ULLR:
                self._encoding_params = torch.tensor(encoding_params, dtype=torch.float32)
                self._coords /= (self._encoding_params - 1).repeat(1, 2)
                self._encoding = BoundingBoxEncoding.NORMALIZED_ULLR
            else:
                raise UnknownBoundingBoxEncodingException()

    def decode(self):
        """Encode the coordinates to BoundingBoxEncoding.ULLR

        """
        if self._encoding == BoundingBoxEncoding.ULLR:
            # bounding box encoding is already in ULLR, no need to decode
            pass
        elif self._encoding == BoundingBoxEncoding.CCWH:
            offsets = self._coords[:, 2:].repeat(1, 2)
            offsets /=2
            offsets *= torch.tensor([-1, -1, 1, 1])
            self._coords = self._coords[:, :2].repeat(1, 2) + offsets
        elif self._encoding == BoundingBoxEncoding.NORMALIZED_ULLR:
            self._coords *= (self._encoding_params - 1).repeat(1, 2)
        else:
            raise UnknownBoundingBoxEncodingException()
        self._encoding = BoundingBoxEncoding.ULLR
        self._encoding_params = None
