"""
Regular, axis-aligned grids.
"""

import numpy as np
cimport numpy as np

class Grids(object):
    """Factory class for creating grids."""

    def __cinit__(self):
        pass

    @classmethod
    def create(cls, dimension, shape, spacing, origin):
        """Returns Grid1D, Grid2D or Grid3D object.

        Args:
            dimension: integer from 1 to 3 (D).
            shape: (D, ) integer array with number of grid points in
                each direction, or single integer if they are the same.
            spacing: (D, ) double array, the distance between
                neighbouring points, or single integer
                if they are the same.
            origin: (D, ) double array, the corner of the grid.

        Returns:
            Grid1D, Grid2D or Grid3D object depending on 'dimension'.
        """

        _shape = cls._expand_int(shape, dimension)
        _spacing = cls._expand(spacing, dimension)
        _origin = cls._expand(origin, dimension)

        grid_class = [Grid1D, Grid2D, Grid3D]

        return grid_class[dimension - 1](_shape, _spacing, _origin)


    @classmethod
    def create_centered(cls, dimension, shape, spacing):
        """Returns correct Grid object with its center at (0, 0, ..., 0).

        Args:
            dimension: integer from 1 to 3 (D).
            shape: (D, ) integer array with number of grid points in
                each direction.
            spacing: (D, ) double array, the distance between
                neighbouring points.

        Returns:
            Grid1D, Grid2D or Grid3D object with origin chosen such that
            its center is at (0, 0, ..., 0).
        """

        origin = -0.5 * (shape - 1.0) * spacing

        return cls.create(dimension, shape, spacing, origin)


    @staticmethod
    def _expand(x, length):
        """Returns (x, x, ..., x) of given length if x is not an array.

        Args:
            x: double or 1D array.
            length: integer indicating length that 'x' should have.

        Returns:
            If 'x' is not an array, return a constant array of length
            'length', else just 'x'.
        """

        if not isinstance(x, np.ndarray):
            return np.ones(length) * x
        else:
            return x

    @staticmethod
    def _expand_int(x, length):
        """Returns (x, x, ..., x) of given length if x is not an array.

        Args:
            x: int or 1D array.
            length: integer indicating length that 'x' should have.

        Returns:
            If 'x' is not an array, return a constant integer array
            of length 'length', else just 'x'.
        """

        if not isinstance(x, np.ndarray):
            return np.ones(length).astype('i') * x
        else:
            return x



    @classmethod
    def load(cls, parent_group, name='grid'):
        """Load grid from group in HDF5 database."""

        group = parent_group[name]
        D = group.attrs['dimension']
        shape = group['shape'][:] #.astype('i')
        spacing = group['spacing'][:]
        origin = group['origin'][:]

        return cls.create(D, shape, spacing, origin)

    @classmethod
    def resize_centered(cls, grid, new_shape):
        """Return new centered grid with given shape by resizing a grid.

        Args:
            new_shape: Array-like of length dimension.
        """

        new_shape_array = np.array(new_shape)
        grid_range = grid.shape * grid.spacing
        new_spacing = grid_range / new_shape_array

        return cls.create_centered(
            grid.dimension, new_shape_array, new_spacing)

    @classmethod
    def create_3D_from_2D(cls, grid2D):
        """Returns a centered 3D grid with the same shape and spacing.

        """

        return cls.create_centered(3, grid2D.shape[0], grid2D.spacing[0])


cdef class Grid:
    """Regular grid of any dimension."""

    def __cinit__(self, shape, spacing, origin):
        """Construct a regular grid.

        Args:
            shape: (D, ) integer array with number of grid points in
                each direction.
            spacing: (D, ) double array, the distance between
                neighbouring points.
            origin: (D, ) double array, the coordinates of the corner
        """

        self._shape = shape.astype('int')
        self._spacing = spacing
        self._origin = origin
        self._points = None
        self._size = np.array(shape).prod()

    @property
    def shape(self):
        return np.array(self._shape)

    @property
    def spacing(self):
        return np.array(self._spacing)

    @property
    def origin(self):
        return np.array(self._origin)

    def set_origin(self, value):
        self._origin = value
        self._points = None

    @property
    def dimension(self):
        return len(self._shape)

    @property
    def size(self):
        return np.prod(self._shape)

    @property
    def points(self):
        """Lazily initialise grid points."""

        if self._points is None:
            print 'Computing grid points.'
            self._points = self._create_points()

        return self._points

    def _create_points(self):
        """Returns grid points."""

        raise NotImplementedError

    def last(self):
        """Returns corner of grid opposite to self.origin."""

        return (self.origin +
                (self.shape - np.ones(self.dimension)) * self.spacing)

    def center(self):
        """Shifts the grid origin so that it is centered."""

        self.set_origin(-0.5 * (self.shape - 1.0) * self.spacing)


    def size(self):
        """Returns the total number of grid points."""

        return self.shape.prod()

    def save(self, parent_group, name='grid'):
        """Save to group in HDF5 database.

        Assume that the group parent_group[name] doesn't exist.

        Args:
            parent_group: Group within which to create the grid group.
            name: Name of the new grid group.
        """

        group = parent_group.create_group(name)
        D = self.dimension
        group.attrs['dimension'] = D
        group.create_dataset('shape', (D,), dtype='int')
        group['shape'][...] = self.shape
        group.create_dataset('spacing', (D,), dtype='float')
        group['spacing'][...] = self.spacing
        group.create_dataset('origin', (D,), dtype='float')
        group['origin'][...] = self.origin

    def copy(self):
        """Return deep copy."""

        return self.__class__(self.shape.copy(), self.spacing.copy(),
                              self.origin.copy())


cdef class Grid1D(Grid):
    """1D regular grid."""

    def _create_points(self):
        """Returns grid points."""

        points = np.linspace(self.origin, self.last(), self.shape)

        return np.array([points]).T


cdef class Grid2D(Grid):
    """2D regular grid."""

    def __cinit__(self, shape, spacing, origin):
        """Extra initialisation for 2D grids."""

        self._shape_0 = shape[0]
        self._shape_1 = shape[1]
        self._spacing_0 = spacing[0]
        self._spacing_1 = spacing[1]
        self._origin_0 = origin[0]
        self._origin_1 = origin[1]

    cdef double to_x(self, int index):
        """Converts the given linear index to x coordinate."""

        return (index / self._shape_1) * self._spacing_0 + self._origin_0

    cdef double to_y(self, int index):
        """Converts the given linear index to y coordinate."""

        return ((index % self._shape_1) * self._spacing_1 +
                self._origin_1)

    def to_frequency_grid(self):
        """Converts this grid to a frequency grid.

        The frequency grid has DC (almost) in the center.

        It might be necessary to first change the units of this grid
        from Angstrom to meters.
        """

        spacing_f = 1.0 / (self.spacing * self.shape)
        origin_f = -(self.shape / 2) * spacing_f

        return Grids.create(2, self.shape, spacing_f, origin_f)

    def to_spatial_grid(self):
        """Converts this frequency grid to a spatial grid.

        Inverse of self.to_frequency_grid.
        """

        spacing_s = 1.0 / (self.spacing * self.shape)
        origin_s = -(self.shape / 2) * spacing_s

        return Grids.create(2, self.shape, spacing_s, origin_s)

    def scale(self, alpha):
        """Scale grid by alpha."""

        return Grids.create(2, self.shape, self.spacing * alpha,
                            self.origin * alpha)


    def _create_points(self):
        """Returns grid points."""

        last = self.last()
        x = [np.linspace(self.origin[i], last[i], num=self.shape[i])
             for i in range(2)]
        X, Y = np.meshgrid(x[0], x[1])
        # Take transposes since x should correspond to rows.
        X = X.T
        Y = Y.T
        cx = X.reshape([X.size,])
        cy = Y.reshape([Y.size,])
        points = (np.vstack((cx, cy))).T # (N, 2) array

        return points


cdef class Grid3D(Grid):
    """3D regular grid."""

    def __cinit__(self, shape, spacing, origin):
        """Extra initialisation for 2D grids."""

        self._shape_0 = shape[0]
        self._shape_1 = shape[1]
        self._shape_2 = shape[2]
        self._spacing_0 = spacing[0]
        self._spacing_1 = spacing[1]
        self._spacing_2 = spacing[2]
        self._origin_0 = origin[0]
        self._origin_1 = origin[1]
        self._origin_2 = origin[2]

    def _create_points(self):
        """Returns grid points."""

        first = self.origin
        last = self.last()

        X, Y, Z = np.mgrid[first[0]:last[0]:self.shape[0]*1j,
                           first[1]:last[1]:self.shape[1]*1j,
                           first[2]:last[2]:self.shape[2]*1j]
        n = X.size
        cx = X.reshape([n,])
        cy = Y.reshape([n,])
        cz = Z.reshape([n,])
        points = (np.vstack((cx, cy, cz))).T

        return points

