# Source code for pointpats.quadrat_statistics

    """
    Quadrat statistics for planar point patterns

    TODO

    - use patch in matplotlib to plot rectangles and hexagons
    - plot chi2 statistics in each cell
    - delete those cells that do not intersect with the window (study area)

    """

    __author__ = "Serge Rey, Wei Kang, Hu Shao"
    __all__ = ["RectangleM", "HexagonM", "QStatistic"]

    from .pointpattern import PointPattern
    import numpy as np
    from matplotlib import pyplot as plt
    import math
    import scipy

    class RectangleM:
        """
        Rectangle grid structure for quadrat-based method.

        Parameters
        ----------
        pp                : :class:`.PointPattern`
                            Point Pattern instance.
        count_column      : integer
                            Number of rectangles in the horizontal
                            direction. Use in pair with count_row to
                            fully specify a rectangle. Incompatible with
                            rectangle_width and rectangle_height.
        count_row         : integer
                            Number of rectangles in the vertical
                            direction. Use in pair with count_column to
                            fully specify a rectangle. Incompatible with
                            rectangle_width and rectangle_height.
        rectangle_width   : float
                            Rectangle width. Use in pair with
                            rectangle_height to fully specify a rectangle.
                            Incompatible with count_column & count_row.
        rectangle_height  : float
                            Rectangle height. Use in pair with
                            rectangle_width to fully specify a rectangle.
                            Incompatible with count_column & count_row.

        Attributes
        ----------
        pp                : :class:`.PointPattern`
                            Point Pattern instance.
        mbb               : array
                            Minimum bounding box for the point pattern.
        points            : array
                            x,y coordinates of the point points.
        count_column      : integer
                            Number of columns.
        count_row         : integer
                            Number of rows.
        num               : integer
                            Number of rectangular quadrats.
        rectangle_width   : float
                            Width of a rectangular quadrat.
        rectangle_height  : float
                            Height of a rectangular quadrat.

        """

        def __init__(
            self, pp, count_column=3, count_row=3, rectangle_width=0, rectangle_height=0
        ):
            self.mbb = pp.mbb
            self.pp = pp
            self.points = np.asarray(pp.points)
            x_range = self.mbb[2] - self.mbb[0]
            y_range = self.mbb[3] - self.mbb[1]
            if rectangle_width and rectangle_height:
                self.rectangle_width = rectangle_width
                self.rectangle_height = rectangle_height

                # calculate column count and row count
                self.count_column = int(math.ceil(x_range / rectangle_width))
                self.count_row = int(math.ceil(y_range / rectangle_height))
            else:
                self.count_column = count_column
                self.count_row = count_row

                # calculate the actual width and height of cell
                self.rectangle_width = x_range / float(count_column)
                self.rectangle_height = y_range / float(count_row)
            self.num = self.count_column * self.count_row

        def point_location_sta(self):
            """
            Count the point events in each cell.

            Returns
            -------
            dict_id_count : dict
                            keys: rectangle id, values: number of point
                            events in each cell.
            """

            dict_id_count = {}
            for i in range(self.count_row):
                for j in range(self.count_column):
                    dict_id_count[j + i * self.count_column] = 0

            for point in self.points:
                index_x = (point[0] - self.mbb[0]) // self.rectangle_width
                index_y = (point[1] - self.mbb[1]) // self.rectangle_height
                if index_x == self.count_column:
                    index_x -= 1
                if index_y == self.count_row:
                    index_y -= 1
                id = index_y * self.count_column + index_x
                dict_id_count[id] += 1
            return dict_id_count

        def plot(self, title="Quadrat Count"):
            """
            Plot rectangle tessellation as well as the number of points falling in each rectangle.

            Parameters
            ----------
            title:   str, optional
                     Title of the plot. Default is "Quadrat Count".

            """

            line_width_cell = 1
            line_color_cell = "red"

            x_min = self.mbb[0]
            y_min = self.mbb[1]

            # draw the point pattern along with its window
            ax = self.pp.plot(window=True, title=title, get_ax=True)

            # draw cells and counts
            x_start_end = [x_min, x_min + self.count_column * self.rectangle_width]
            for row in range(self.count_row + 1):
                y = y_min + row * self.rectangle_height
                ax.plot(x_start_end, [y, y], lw=line_width_cell, color=line_color_cell)
            y_start_end = [y_min, y_min + self.count_row * self.rectangle_height]
            for column in range(self.count_column + 1):
                x = x_min + column * self.rectangle_width
                ax.plot([x, x], y_start_end, lw=line_width_cell, color=line_color_cell)

            dict_id_count = self.point_location_sta()
            for x in range(self.count_column):
                for y in range(self.count_row):
                    cell_id = x + y * self.count_column
                    count = dict_id_count[cell_id]
                    position_x = x_min + self.rectangle_width * (x + 0.5)
                    position_y = y_min + self.rectangle_height * (y + 0.5)
                    ax.text(position_x, position_y, str(count), ha="center", va="center")
            return ax

    class HexagonM:
        """
        Hexagon grid structure for quadrat-based method.

        Parameters
        ----------
        pp                : :class:`.PointPattern`
                            Point Pattern instance.
        lh                : float
                            Hexagon length (hexagon).

        Attributes
        ----------
        pp                : :class:`.PointPattern`
                            Point Pattern instance.
        h_length          : float
                            Hexagon length (hexagon).
        mbb               : array
                            Minimum bounding box for the point pattern.
        points            : array
                            x,y coordinates of the point points.
        h_length          : float
                            Hexagon length (hexagon).
        count_row_even    : integer
                            Number of even rows.
        count_row_odd     : integer
                            Number of odd rows.
        count_column      : integer
                            Number of columns.
        num               : integer
                            Number of hexagonal quadrats.

        """

        def __init__(self, pp, lh):

            self.points = np.asarray(pp.points)
            self.pp = pp
            self.h_length = lh
            self.mbb = pp.mbb
            range_x = self.mbb[2] - self.mbb[0]
            range_y = self.mbb[3] - self.mbb[1]

            # calculate column count
            self.count_column = 1
            if self.h_length / 2.0 < range_x:
                temp = math.ceil((range_x - self.h_length / 2) / (1.5 * self.h_length))
                self.count_column += int(temp)

            # calculate row count for the even columns
            self.semi_height = self.h_length * math.cos(math.pi / 6)
            self.count_row_even = 1
            if self.semi_height < range_y:
                temp = math.ceil((range_y - self.semi_height) / (self.semi_height * 2))
                self.count_row_even += int(temp)

            # for the odd columns
            self.count_row_odd = int(math.ceil(range_y / (self.semi_height * 2)))

            # quadrat number
            self.num = self.count_row_odd * (
                (self.count_column // 2) + self.count_column % 2
            ) + self.count_row_even * (self.count_column // 2)

        def point_location_sta(self):
            """
            Count the point events in each hexagon cell.

            Returns
            -------
            dict_id_count : dict
                            keys: rectangle id, values: number of point
                            events in each hexagon cell.
            """
            semi_cell_length = self.h_length / 2.0
            dict_id_count = {}

            # even row may be equal with odd row or 1 more than odd row
            for i in range(self.count_row_even):
                for j in range(self.count_column):
                    if (
                        self.count_row_even != self.count_row_odd
                        and i == self.count_row_even - 1
                    ):
                        if j % 2 == 1:
                            continue
                    dict_id_count[j + i * self.count_column] = 0

            x_min = self.mbb[0]
            y_min = self.mbb[1]
            x_max = self.mbb[2]
            y_max = self.mbb[3]
            points = np.array(self.points)
            for point in points:
                # find the possible x index
                intercept_degree_x = (point[0] - x_min) // semi_cell_length

                # find the possible y index
                possible_y_index_even = int(
                    (point[1] + self.semi_height - y_min) / (self.semi_height * 2)
                )
                possible_y_index_odd = int((point[1] - y_min) / (self.semi_height * 2))
                if intercept_degree_x % 3 != 1:
                    center_index_x = (intercept_degree_x + 1) // 3
                    center_index_y = possible_y_index_odd
                    if center_index_x % 2 == 0:
                        center_index_y = possible_y_index_even
                    dict_id_count[center_index_x + center_index_y * self.count_column] += 1
                else:  # two columns of cells can be possible
                    center_index_x = intercept_degree_x // 3
                    center_x = center_index_x * semi_cell_length * 3 + x_min
                    center_index_y = possible_y_index_odd
                    center_y = (center_index_y * 2 + 1) * self.semi_height + y_min
                    if center_index_x % 2 == 0:
                        center_index_y = possible_y_index_even
                        center_y = center_index_y * self.semi_height * 2 + y_min

                    if point[1] > center_y:  # compare the upper bound
                        x0 = center_x + self.h_length
                        y0 = center_y
                        x1 = center_x + semi_cell_length
                        y1 = center_y + self.semi_height
                        indicator = -(
                            point[1]
                            - (
                                (y0 - y1) / (x0 - x1) * point[0]
                                + (x0 * y1 - x1 * y0) / (x0 - x1)
                            )
                        )
                    else:  # compare the lower bound
                        x0 = center_x + semi_cell_length
                        y0 = center_y - self.semi_height
                        x1 = center_x + self.h_length
                        y1 = center_y
                        indicator = point[1] - (
                            (y0 - y1) / (x0 - x1) * point[0]
                            + (x0 * y1 - x1 * y0) / (x0 - x1)
                        )
                    if indicator <= 0:
                        # we select right hexagon instead of the left
                        center_index_x += 1
                        center_index_y = possible_y_index_odd
                        if center_index_x % 2 == 0:
                            center_index_y = possible_y_index_even
                    dict_id_count[center_index_x + center_index_y * self.count_column] += 1
            return dict_id_count

        def plot(self, title="Quadrat Count"):
            """
            Plot hexagon quadrats as well as the number of points falling in each quadrat.

            Parameters
            ----------
            title:   str, optional
                     Title of the plot. Default is "Quadrat Count".

            """
            line_width_cell = 1
            line_color_cell = "red"

            # draw the point pattern along with its window
            ax = self.pp.plot(window=True, title=title, get_ax=True)

            x_min = self.mbb[0]
            y_min = self.mbb[1]

            # draw cells and counts
            dict_id_count = self.point_location_sta()
            for id in dict_id_count.keys():
                index_x = id % self.count_column
                index_y = id // self.count_column
                center_x = index_x * self.h_length / 2.0 * 3.0 + x_min
                center_y = index_y * self.semi_height * 2.0 + y_min
                if index_x % 2 == 1:  # for the odd columns
                    center_y = (index_y * 2.0 + 1) * self.semi_height + y_min
                list_points_cell = []
                list_points_cell.append([center_x + self.h_length, center_y])
                list_points_cell.append(
                    [center_x + self.h_length / 2, center_y + self.semi_height]
                )
                list_points_cell.append(
                    [center_x - self.h_length / 2, center_y + self.semi_height]
                )
                list_points_cell.append([center_x - self.h_length, center_y])
                list_points_cell.append(
                    [center_x - self.h_length / 2, center_y - self.semi_height]
                )
                list_points_cell.append(
                    [center_x + self.h_length / 2, center_y - self.semi_height]
                )
                list_points_cell.append([center_x + self.h_length, center_y])
                ax.plot(
                    np.array(list_points_cell)[:, 0],
                    np.array(list_points_cell)[:, 1],
                    lw=line_width_cell,
                    color=line_color_cell,
                )

                ax.text(center_x, center_y, str(dict_id_count[id]), ha="center", va="center")
            return ax

    class QStatistic:
        """
        Quadrat analysis of point pattern.

        Parameters
        ----------
        pp                : :class:`.PointPattern` or numpy.ndarray
                            Point Pattern instance, or (n_observations, 2) array
                            that can be used to construct a Point Pattern instance.
        shape             : string
                            Grid structure. Either "rectangle" or "hexagon".
                            Default is "rectangle".
        nx                : integer
                            Number of rectangles in the horizontal
                            direction. Only when shape is specified as
                            "rectangle" will nx be considered.
        ny                : integer
                            Number of rectangles in the vertical direction.
                            Only when shape is specified as "rectangle"
                            will ny be considered.
        rectangle_width   : float
                            Rectangle width. Use in pair with
                            rectangle_height to fully specify a rectangle.
                            Incompatible with nx & ny.
        rectangle_height  : float
                            Rectangle height. Use in pair with
                            rectangle_width to fully specify a rectangle.
                            Incompatible with nx & ny.
        lh                : float
                            Hexagon length (hexagon). Only when shape is
                            specified as "hexagon" will lh be considered.
                            Incompatible with nx & ny.
        realizations      : :class:`PointProcess`
                            Point process instance with more than 1 point
                            pattern realizations which would be used for
                            simulation based inference. Default is 0
                            where no simulation based inference is
                            performed.

        Attributes
        ----------
        pp                : :class:`.PointPattern`
                            Point Pattern instance.
        mr                : :class:`.RectangleM` or :class:`.HexagonM`
                            RectangleM or HexagonM instance.
        chi2              : float
                            Chi-squared test statistic for the observed
                            point pattern pp.
        df                : integer
                            Degree of freedom.
        chi2_pvalue       : float
                            p-value based on analytical chi-squared
                            distribution.
        chi2_r_pvalue     : float
                            p-value based on simulated sampling
                            distribution. Only available when
                            realizations is correctly specified.
        chi2_realizations : array
                            Chi-squared test statistics calculated for
                            all the simulated csr point patterns.
        """

        def __init__(self, pp, shape="rectangle", nx=3, ny=3, rectangle_width=0,
                     rectangle_height=0, lh=10, realizations=0):
            if isinstance(pp, np.ndarray):
                pp = PointPattern(pp)
            self.pp = pp
            if shape == "rectangle":
                self.mr = RectangleM(pp, count_column=nx, count_row=ny,
                                     rectangle_width=rectangle_width,
                                     rectangle_height = rectangle_height)
            elif shape == "hexagon":
                self.mr = HexagonM(pp, lh)
            else:
                raise ValueError(
                    f'shape type {shape} not understood. Must be either "rectangle" or'
                    ' "hexagon"'
                )

            # calculate chi2 test statisitc for the observed point pattern
            dict_id_count = self.mr.point_location_sta()
            self.chi2, self.chi2_pvalue = scipy.stats.chisquare(
                list(dict_id_count.values())
            )

            self.df = self.mr.num - 1

            # when realizations is specified, perform simulation based
            # inference.
            if realizations:
                reals = realizations.realizations
                sim_n = realizations.samples
                chi2_realizations = []  # store test statistics for all the
                # similations

                for i in range(sim_n):
                    if shape == "rectangle":
                        mr_temp = RectangleM(reals[i], count_column=nx, count_row=ny)
                    elif shape == "hexagon":
                        mr_temp = HexagonM(reals[i], lh)
                    id_count_temp = mr_temp.point_location_sta().values()

                    # calculate test statistics for simulated point patterns
                    chi2_sim, p = scipy.stats.chisquare(list(id_count_temp))
                    chi2_realizations.append(chi2_sim)
                self.chi2_realizations = np.array(chi2_realizations)

                # calculate pseudo pvalue
                above_chi2 = self.chi2_realizations >= self.chi2
                larger_chi2 = sum(above_chi2)
                self.chi2_r_pvalue = (larger_chi2 + 1.0) / (sim_n + 1.0)

        def plot(self, title="Quadrat Count"):
            """
            Plot quadrats as well as the number of points falling in each quadrat.

            Parameters
            ----------
            title:   str, optional
                     Title of the plot. Default is "Quadrat Count".

            """

            return self.mr.plot(title=title)