import numpy as np
import scipy.stats as stats

from UQpy.SampleMethods.Strata.strata import Strata


class VoronoiStrata(Strata):
    """
    Define a geometric decomposition of the n-dimensional unit hypercube into disjoint and space-filling
    Voronoi strata.

    ``VoronoiStrata`` is a child class of the ``Strata`` class.

    **Inputs:**

    * **seeds** (`ndarray`):
        An array of dimension `N x n` specifying the seeds of all strata. The seeds of the strata are the
        coordinates of the point inside each stratum that defines the stratum.

        The user must provide `seeds` or `nseeds` and `dimension`

    * **nseeds** (`int`):
        The number of seeds to randomly generate. Seeds are generated by random sampling on the unit hypercube.

        The user must provide `seeds` or `nseeds` and `dimension`

    * **dimension** (`ndarray`):
        The dimension of the unit hypercube in which to generate random seeds. Used only if `nseeds` is provided.

        The user must provide `seeds` or `nseeds` and `dimension`

    * **niters** (`int`)
        Number of iterations to perform to create a Centroidal Voronoi decomposition.

        If `niters = 0`, the Voronoi decomposition is based on the provided or generated seeds.

        If :math:`niters \ge 1`, the seed points are moved to the centroids of the Voronoi cells in each iteration and
        the a new Voronoi decomposition is performed. This process is repeated `niters` times to create a Centroidal
        Voronoi decomposition.

    * **random_state** (None or `int` or ``numpy.random.RandomState`` object):
        Random seed used to initialize the pseudo-random number generator. Default is None.

        If an integer is provided, this sets the seed for an object of ``numpy.random.RandomState``. Otherwise, the
        object itself can be passed directly.

    * **verbose** (`Boolean`):
        A boolean declaring whether to write text to the terminal.


    **Attributes:**

    * **seeds** (`ndarray`):
        An array of dimension `N x n` containing the seeds of all strata. The seeds of the strata are the
        coordinates of the point inside each stratum that defines the stratum.

        If :math:`niters > 1` the `seeds` attribute will differ from the `seeds` input due to the iterations.

    * **vertices** (`list`)
        A list of the vertices for each Voronoi stratum on the unit hypercube.

    * **voronoi** (`object` of ``scipy.spatial.Voronoi``)
        Defines a Voronoi decomposition of the set of reflected points. When creating the Voronoi decomposition on
        the unit hypercube, the code reflects the points on the unit hypercube across all faces of the unit hypercube.
        This causes the Voronoi decomposition to create edges along the faces of the hypercube.

        This object is not the Voronoi decomposition of the unit hypercube. It is the Voronoi decomposition of all
        points and their reflections from which the unit hypercube is extracted.

        To access the vertices in the unit hypercube, see the attribute `vertices`.

    * **volume** (`ndarray`):
        An array of dimension `(nstrata, )` containing the volume of each Voronoi stratum in the unit hypercube.

    **Methods:**
    """

    def __init__(self, seeds=None, nseeds=None, dimension=None, niters=1, random_state=None, verbose=False):
        super().__init__(random_state=random_state, seeds=seeds, verbose=verbose)

        self.nseeds = nseeds
        self.dimension = dimension
        self.niters = niters
        self.voronoi = None
        self.vertices = []

        if self.seeds is not None:
            if self.nseeds is not None or self.dimension is not None:
                print("UQpy: Ignoring 'nseeds' and 'dimension' attributes because 'seeds' are provided")
            self.dimension = self.seeds.shape[1]

        self.stratify()

    def stratify(self):
        """
        Performs the Voronoi stratification.
        """
        if self.verbose:
            print('UQpy: Creating Voronoi stratification ...')

        initial_seeds = self.seeds
        if self.seeds is None:
            initial_seeds = stats.uniform.rvs(size=[self.nseeds, self.dimension], random_state=self.random_state)

        if self.niters == 0:
            self.voronoi, bounded_regions = self.voronoi_unit_hypercube(initial_seeds)

            cent, vol = [], []
            for region in bounded_regions:
                vertices = self.voronoi.vertices[region + [region[0]], :]
                centroid, volume = self.compute_voronoi_centroid_volume(vertices)
                self.vertices.append(vertices)
                cent.append(centroid[0, :])
                vol.append(volume)

            self.volume = np.asarray(vol)
        else:
            for i in range(self.niters):
                self.voronoi, bounded_regions = self.voronoi_unit_hypercube(initial_seeds)

                cent, vol = [], []
                for region in bounded_regions:
                    vertices = self.voronoi.vertices[region + [region[0]], :]
                    centroid, volume = self.compute_voronoi_centroid_volume(vertices)
                    self.vertices.append(vertices)
                    cent.append(centroid[0, :])
                    vol.append(volume)

                initial_seeds = np.asarray(cent)
                self.volume = np.asarray(vol)

        self.seeds = initial_seeds

        if self.verbose:
            print('UQpy: Voronoi stratification created.')

    @staticmethod
    def voronoi_unit_hypercube(seeds):
        """
        This function reflects the seeds across all faces of the unit hypercube and creates a Voronoi decomposition of
        using all the points and their reflections. This allows a Voronoi decomposition that is bounded on the unit
        hypercube to be extracted.

        **Inputs:**

        * **seeds** (`ndarray`):
            Coordinates of points in the unit hypercube from which to define the Voronoi decomposition.

        **Output/Returns:**

        * **vor** (``scipy.spatial.Voronoi`` object):
            Voronoi decomposition of the complete set of points and their reflections.

        * **bounded_regions** (see `regions` attribute of ``scipy.spatial.Voronoi``)
            Indices of the Voronoi vertices forming each Voronoi region for those regions lying inside the unit
            hypercube.
        """

        from scipy.spatial import Voronoi

        # Mirror the seeds in both low and high directions for each dimension
        bounded_points = seeds
        dimension = seeds.shape[1]
        for i in range(dimension):
            seeds_del = np.delete(bounded_points, i, 1)
            if i == 0:
                points_temp1 = np.hstack([np.atleast_2d(-bounded_points[:, i]).T, seeds_del])
                points_temp2 = np.hstack([np.atleast_2d(2 - bounded_points[:, i]).T, seeds_del])
            elif i == dimension - 1:
                points_temp1 = np.hstack([seeds_del, np.atleast_2d(-bounded_points[:, i]).T])
                points_temp2 = np.hstack([seeds_del, np.atleast_2d(2 - bounded_points[:, i]).T])
            else:
                points_temp1 = np.hstack(
                    [seeds_del[:, :i], np.atleast_2d(-bounded_points[:, i]).T, seeds_del[:, i:]])
                points_temp2 = np.hstack(
                    [seeds_del[:, :i], np.atleast_2d(2 - bounded_points[:, i]).T, seeds_del[:, i:]])
            seeds = np.append(seeds, points_temp1, axis=0)
            seeds = np.append(seeds, points_temp2, axis=0)

        vor = Voronoi(seeds, incremental=True)

        regions = [None] * bounded_points.shape[0]

        for i in range(bounded_points.shape[0]):
            regions[i] = vor.regions[vor.point_region[i]]

        bounded_regions = regions

        return vor, bounded_regions

    @staticmethod
    def compute_voronoi_centroid_volume(vertices):
        """
        This function computes the centroid and volume of a Voronoi cell from its vertices.

        **Inputs:**

        * **vertices** (`ndarray`):
            Coordinates of the vertices that define the Voronoi cell.

        **Output/Returns:**

        * **centroid** (`ndarray`):
            Centroid of the Voronoi cell.

        * **volume** (`ndarray`):
            Volume of the Voronoi cell.
        """

        from scipy.spatial import Delaunay, ConvexHull

        tess = Delaunay(vertices)
        dimension = np.shape(vertices)[1]

        w = np.zeros((tess.nsimplex, 1))
        cent = np.zeros((tess.nsimplex, dimension))
        for i in range(tess.nsimplex):
            # pylint: disable=E1136
            ch = ConvexHull(tess.points[tess.simplices[i]])
            w[i] = ch.volume
            cent[i, :] = np.mean(tess.points[tess.simplices[i]], axis=0)

        volume = np.sum(w)
        centroid = np.matmul(np.divide(w, volume).T, cent)

        return centroid, volume
