'''
This module defines the :class:`Cuboid` class, which represents an n-dimensional
cuboid that is aligned with the axes of a Cartesian coordinate system.

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
'''

import itertools
from typing import Tuple

import numpy as np



class Cuboid:
    """ class that represents a cuboid in :math:`n` dimensions """
    
    
    def __init__(self, pos, size, mutable: bool = True):
        """ defines a cuboid from a position and a size vector 

        Args:
            pos (list):
                The position of the lower left corner. The length of this list
                determines the dimensionality of space
            size (list):
                The size of the cuboid along each dimension.
            mutable (bool):
                Flag determining whether the cuboid parameters can be changed
        """
        self._mutable = mutable
        # set position and adjust mutable status later
        self.pos = np.array(pos, copy=True)
        self.size = size  # implicitly sets correct shape
        self.pos.flags.writeable = self.mutable 
        
        
    @property
    def size(self):
        return self._size
    
    @size.setter
    def size(self, value):
        self._size = np.array(value, self.pos.dtype)  # make copy
        if self.pos.shape != self._size.shape:
            raise ValueError('Size vector (dim=%d) must have the same '
                             'dimension as the position vector (dim=%d)' %
                             (len(self._size), len(self.pos)))
        
        # flip Cuboid with negative size
        neg = (self._size < 0)
        self.pos[neg] += self._size[neg]
        self._size = np.abs(self._size)
        self._size.flags.writeable = self.mutable
        

    @property
    def corners(self):
        """ return coordinates of two extreme corners defining the cuboid """
        return np.copy(self.pos), self.pos + self.size
    
            
    @property
    def mutable(self) -> bool:
        return self._mutable
    
    @mutable.setter
    def mutable(self, value: bool):
        self._mutable = bool(value)
        self.pos.flags.writeable = self._mutable
        self._size.flags.writeable = self._mutable

        
    @classmethod
    def from_points(cls, p1, p2, **kwargs) -> "Cuboid":
        """ create cuboid from two points
        
        Args:
            p1 (list): Coordinates of first corner point
            p2 (list): Coordinates of second corner point
            
        Returns:
            Cuboid: cuboid with positive size
        """
        p1 = np.asarray(p1)
        p2 = np.asarray(p2)
        return cls(p1, p2 - p1, **kwargs)
    
    
    @classmethod
    def from_bounds(cls, bounds, **kwargs) -> "Cuboid":
        """ create cuboid from bounds
        
        Args:
            bounds (list): Two dimensional array of axes bounds
            
        Returns:
            Cuboid: cuboid with positive size
        """
        bounds = np.asarray(bounds).reshape(-1, 2)
        return cls(bounds[:, 0], bounds[:, 1] - bounds[:, 0], **kwargs)
    
    
    @classmethod
    def from_centerpoint(cls, centerpoint, size, **kwargs) \
            -> "Cuboid":
        """ create cuboid from two points
        
        Args:
            centerpoint (list): Coordinates of the center
            size (list): Size of the cuboid
            
        Returns:
            Cuboid: cuboid with positive size
        """        
        centerpoint = np.asarray(centerpoint)
        size = np.asarray(size)
        return cls(centerpoint - size / 2, size, **kwargs)
    
    
    def copy(self) -> "Cuboid":
        return self.__class__(self.pos, self.size)
        
        
    def __repr__(self):
        return "{cls}(pos={pos}, size={size})".format(
                        cls=self.__class__.__name__,
                        pos=self.pos, size=self.size)
            
            
    def __add__(self, other: "Cuboid"):
        """ The sum of two cuboids is the minimal cuboid enclosing both """
        if isinstance(other, Cuboid):
            if self.dim != other.dim:
                raise RuntimeError('Incompatible dimensions')
            a1, a2 = self.corners
            b1, b2 = other.corners
            return self.__class__.from_points(np.minimum(a1, b1),
                                              np.maximum(a2, b2))
            
        else:
            raise NotImplementedError('Cannot add cuboid and '
                                      f'{other.__class__.__name__}')
            
            
    def __eq__(self, other):
        """ override the default equality test """
        if not isinstance(other, self.__class__):
            return NotImplemented
        return (np.all(self.pos == other.pos) and 
                np.all(self.size == other.size))

    
    @property
    def dim(self) -> int:
        return len(self.pos)
            

    @property
    def bounds(self) -> Tuple[Tuple[float, float], ...]:
        return tuple((p, p + s) for p, s in zip(self.pos, self.size))
        

    @property
    def vertices(self):
        """ return the coordinates of all the corners """
        return list(itertools.product(*self.bounds))


    @property
    def diagonal(self) -> float:
        """ returns the length of the diagonal """
        return np.linalg.norm(self.size)  # type: ignore
    
    
    @property
    def surface_area(self) -> float:
        """ surface area of a cuboid in :math:`n` dimensions.
        
        The surface area is the volume of the (:math:`n-1`)-dimensional
        hypercubes that bound the current cuboid:
        
            * :math:`n=1`: the number of end points (2)
            * :math:`n=2`: the perimeter of the rectangle
            * :math:`n=3`: the surface area of the cuboid
        """
        sides = self.size
        null = (sides == 0)
        null_count = null.sum()
        if null_count == 0:
            return 2 * np.sum(np.product(sides) / sides)  # type: ignore
        elif null_count == 1:
            return 2 * np.product(sides[~null])  # type: ignore
        else:
            return 0


    @property
    def centroid(self):
        return self.pos + self.size / 2
    
    @centroid.setter
    def centroid(self, center):
        self.pos[:] = np.asanyarray(center) - self.size / 2
    
    
    @property
    def volume(self) -> float:
        return np.prod(self.size)  # type: ignore
    

    def buffer(self, amount=0, inplace=False):
        """ dilate the cuboid by a certain amount in all directions """
        amount = np.asarray(amount)
        if inplace:
            self.pos -= amount
            self.size += 2*amount
            return self
        else:
            return self.__class__(self.pos - amount, self.size + 2*amount)
    
        
    def contains_point(self, points):
        """ returns a True when `points` are within the Cuboid
        
        Args:
            points (list): List of point coordinates
            
        Returns:
            np.ndarray: list of booleans indicating which points are inside
        """
        points = np.asarray(points)
        if len(points) == 0:
            return points
        
        if points.shape[-1] != self.dim:
            raise ValueError('Last dimension of `points` must agree with '
                             'cuboid dimension %d' % self.dim)
        
        c1, c2 = self.corners
        return np.all(c1 <= points, axis=-1) & np.all(points <= c2, axis=-1)
        


def asanyarray_flags(data: np.ndarray, dtype=None, writeable: bool = True):
    """ turns data into an array and sets the respective flags.
    
    A copy is only made if necessary
    
    Args:
        data (:class:`numpy.ndarray`): numpy array that whose flags are adjusted
        dtype: the resulant dtype
        writeable (bool): Flag determining whether the results is writable
        
    Returns:
        :class:`numpy.ndarray`:
            array with same data as `data` but with flags adjusted.
    """
    try:
        data_writeable = data.flags.writeable
    except AttributeError:
        # `data` did not have the writeable flag => it's not a numpy array  
        result = np.array(data, dtype)
    else:
        if data_writeable != writeable:
            # need to make a copy because the flags differ
            result = np.array(data, dtype)
        else:
            # might have to make a copy to adjust the dtype
            result = np.asanyarray(data, dtype)
            
    # set the flags and return the array
    result.flags.writeable = writeable
    return result
    