# Source code for pointpats.quadrat_statistics

    """
    Quadrat statistics for planar point patterns

    - First argument everywhere is a NumPy ndarray of point coordinates (n, 2).
    - MBB is computed internally from the array.
    - Optional `window` (shapely geometry). If omitted, uses the MBB rectangle.
    - Optional `rng` for reproducible simulation-based inference.
    - Rectangle and hexagon grids supported.
    - Plotting is self-contained (matplotlib only): scatters points, draws window, draws grid,
      annotates counts, and optionally annotates per-cell chi-square contributions.

    """

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

    import math

    import geopandas as gpd
    import numpy as np
    import scipy.stats
    from matplotlib import pyplot as plt
    from matplotlib.collections import PatchCollection
    from matplotlib.patches import Polygon as MplPolygon
    from shapely.geometry import MultiPolygon, Polygon, box

    from .random import poisson

    # -----------------------------------------------------------------------------
    # Helpers
    # -----------------------------------------------------------------------------
    def _as_points_array(points) -> np.ndarray:
        """
        Normalize input into an (n, 2) float ndarray.

        Accepts:
          - array-like shaped (n, 2)
          - GeoPandas GeoSeries / GeometryArray of Points (and optionally MultiPoints)
          - PointPattern

        Returns:
          - np.ndarray of shape (n, 2)

        Raises:
          - ValueError for empty inputs or non-2D coordinate inputs
          - TypeError for unsupported geometry types
        """
        if isinstance(points, gpd.GeoDataFrame):
            points = points.geometry
        if isinstance(points, gpd.GeoSeries):
            points = points.dropna()
            if not points.geom_type.isin(["Point", "MultiPoint"]).all():
                raise TypeError(
                    "GeoSeries must contain Point geometries (optionally MultiPoint). "
                    f"Got geometry type: {getattr(points, 'geom_type', type(points).__name__)}"
                )

            coords = points.get_coordinates()

            if len(coords) == 0:
                raise ValueError("no non-empty Point geometries.")

            return coords.values.astype(float)

        # ---- PointPattern ---
        if hasattr(points, "points"):
            return points.points.values

        # ---- Generic array-like path ----
        pts = np.asarray(points)
        if pts.ndim != 2 or pts.shape[1] != 2:
            raise ValueError(f"points must be a (n, 2) array-like. Got shape {pts.shape}.")
        if pts.size == 0:
            raise ValueError("points is empty; cannot compute mbb.")
        return pts

    def _compute_mbb(points: np.ndarray) -> np.ndarray:
        x_min = float(np.min(points[:, 0]))
        y_min = float(np.min(points[:, 1]))
        x_max = float(np.max(points[:, 0]))
        y_max = float(np.max(points[:, 1]))
        return np.array([x_min, y_min, x_max, y_max], dtype=float)

    def _ensure_window(window, mbb):
        if window is None:
            return box(mbb[0], mbb[1], mbb[2], mbb[3])
        return window

    def _coerce_rng(rng):
        """Coerce rng input to a NumPy Generator.

        Parameters
        ----------
        rng : None | int | numpy.random.Generator | numpy.random.RandomState | numpy.random.BitGenerator
            Random number generator specification.

        Returns
        -------
        numpy.random.Generator
        """
        if rng is None:
            return np.random.default_rng()

        if isinstance(rng, np.random.Generator):
            return rng

        # Accept legacy RandomState by deriving a seed from it.
        if isinstance(rng, np.random.RandomState):
            # Draw a seed deterministically from the RandomState stream.
            seed = int(rng.randint(0, 2**32 - 1, dtype=np.uint32))
            return np.random.default_rng(seed)

        # If a BitGenerator is passed, wrap it.
        if isinstance(rng, np.random.BitGenerator):
            return np.random.Generator(rng)

        # Int-like seeds
        if isinstance(rng, (int, np.integer)):
            return np.random.default_rng(int(rng))

        # Fall back to NumPy’s coercion (e.g., SeedSequence)
        return np.random.default_rng(rng)

    def _window_to_paths(window_geom):
        """Return list of (x, y) arrays for plotting boundary(ies)."""
        paths = []
        if window_geom is None:
            return paths

        if isinstance(window_geom, Polygon):
            x, y = window_geom.exterior.xy
            paths.append((np.asarray(x), np.asarray(y)))
            for ring in window_geom.interiors:
                x, y = ring.xy
                paths.append((np.asarray(x), np.asarray(y)))
            return paths

        if isinstance(window_geom, MultiPolygon):
            for poly in window_geom.geoms:
                paths.extend(_window_to_paths(poly))
            return paths

        if hasattr(window_geom, "geoms"):
            for g in window_geom.geoms:
                paths.extend(_window_to_paths(g))
        return paths

    def _scatter_points(ax, points):
        ax.scatter(points[:, 0], points[:, 1], s=10)

    def _draw_window(ax, window_geom):
        for x, y in _window_to_paths(window_geom):
            ax.plot(x, y, lw=1.5)

    def _cell_center_from_bounds(x0, y0, x1, y1):
        return (0.5 * (x0 + x1), 0.5 * (y0 + y1))

    # -----------------------------------------------------------------------------
    # Rectangle grid
    # -----------------------------------------------------------------------------

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

        def __init__(
            self,
            pp,
            count_column=3,
            count_row=3,
            rectangle_width=0,
            rectangle_height=0,
            window=None,
        ):
            self.points = _as_points_array(pp)
            self.mbb = _compute_mbb(self.points)
            self.window = _ensure_window(window, self.mbb)

            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 = float(rectangle_width)
                self.rectangle_height = float(rectangle_height)
                self.count_column = int(math.ceil(x_range / self.rectangle_width))
                self.count_row = int(math.ceil(y_range / self.rectangle_height))
            else:
                self.count_column = int(count_column)
                self.count_row = int(count_row)
                self.rectangle_width = x_range / float(self.count_column)
                self.rectangle_height = y_range / float(self.count_row)

            self.num = self.count_column * self.count_row

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

            x_min, y_min = self.mbb[0], self.mbb[1]
            for point in self.points:
                index_x = (point[0] - x_min) // self.rectangle_width
                index_y = (point[1] - y_min) // self.rectangle_height

                if index_x == self.count_column:
                    index_x -= 1
                if index_y == self.count_row:
                    index_y -= 1

                cell_id = int(index_y) * self.count_column + int(index_x)
                dict_id_count[cell_id] += 1

            return dict_id_count

        def _rect_cell_bounds(self, ix, iy):
            x0 = self.mbb[0] + ix * self.rectangle_width
            y0 = self.mbb[1] + iy * self.rectangle_height
            x1 = x0 + self.rectangle_width
            y1 = y0 + self.rectangle_height
            return x0, y0, x1, y1

        def _rect_cell_polygon(self, ix, iy):
            x0, y0, x1, y1 = self._rect_cell_bounds(ix, iy)
            return box(x0, y0, x1, y1)

        def _iter_cells(self):
            for iy in range(self.count_row):
                for ix in range(self.count_column):
                    cell_id = ix + iy * self.count_column
                    poly = self._rect_cell_polygon(ix, iy)
                    x0, y0, x1, y1 = self._rect_cell_bounds(ix, iy)
                    cx, cy = _cell_center_from_bounds(x0, y0, x1, y1)
                    yield cell_id, poly, (cx, cy)

        def plot(
            self,
            title="Quadrat Count",
            show="counts",
            chi2_contrib=None,
        ):
            fig, ax = plt.subplots()
            ax.set_title(title)

            _scatter_points(ax, self.points)
            _draw_window(ax, self.window)

            dict_id_count = self.point_location_sta()

            patches = []
            ann_xy = []
            cell_ids = []

            for cell_id, poly, (cx, cy) in self._iter_cells():
                patches.append(MplPolygon(np.asarray(poly.exterior.coords), closed=True))
                ann_xy.append((cx, cy))
                cell_ids.append(cell_id)

            pc = PatchCollection(patches, linewidths=1.0, edgecolor="red", facecolor="none")
            ax.add_collection(pc)

            if show == "counts":
                for (cx, cy), cell_id in zip(ann_xy, cell_ids):
                    ax.text(
                        cx, cy, str(dict_id_count.get(cell_id, 0)), ha="center", va="center"
                    )

            elif show == "chi2":
                if chi2_contrib is None:
                    raise ValueError('chi2_contrib must be provided when show="chi2".')
                chi2_contrib = np.asarray(chi2_contrib, dtype=float)
                if chi2_contrib.shape[0] != len(cell_ids):
                    raise ValueError(
                        "chi2_contrib length must match number of plotted cells "
                        f"({len(cell_ids)}). Got {chi2_contrib.shape[0]}."
                    )
                pc.set_array(chi2_contrib)
                pc.set_alpha(0.6)
                pc.set_facecolor(None)
                plt.colorbar(pc, ax=ax, label="Chi-square contribution")

                for (cx, cy), v in zip(ann_xy, chi2_contrib):
                    ax.text(cx, cy, f"{v:.2f}", ha="center", va="center")

            else:
                raise ValueError('show must be "counts" or "chi2".')

            ax.set_aspect("equal", adjustable="box")
            return ax, cell_ids

    # -----------------------------------------------------------------------------
    # Hexagon grid
    # -----------------------------------------------------------------------------

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

        def __init__(self, pp, lh, window=None):
            self.points = _as_points_array(pp)
            self.h_length = float(lh)
            self.mbb = _compute_mbb(self.points)
            self.window = _ensure_window(window, self.mbb)

            range_x = self.mbb[2] - self.mbb[0]
            range_y = self.mbb[3] - self.mbb[1]

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

            self.semi_height = self.h_length * math.cos(math.pi / 6.0)
            self.count_row_even = 1
            if self.semi_height < range_y:
                temp = math.ceil((range_y - self.semi_height) / (2.0 * self.semi_height))
                self.count_row_even += int(temp)

            self.count_row_odd = int(math.ceil(range_y / (2.0 * self.semi_height)))

            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):
            semi_cell_length = self.h_length / 2.0
            dict_id_count = {}

            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
                        and j % 2 == 1
                    ):
                        continue
                    dict_id_count[j + i * self.count_column] = 0

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

            for point in self.points:
                intercept_degree_x = (point[0] - x_min) // semi_cell_length

                possible_y_index_even = int(
                    (point[1] + self.semi_height - y_min) / (2.0 * self.semi_height)
                )
                possible_y_index_odd = int((point[1] - y_min) / (2.0 * self.semi_height))

                if intercept_degree_x % 3 != 1:
                    center_index_x = int((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:
                    center_index_x = int(intercept_degree_x // 3)
                    center_x = center_index_x * semi_cell_length * 3.0 + x_min

                    center_index_y = possible_y_index_odd
                    center_y = (center_index_y * 2.0 + 1.0) * self.semi_height + y_min

                    if center_index_x % 2 == 0:
                        center_index_y = possible_y_index_even
                        center_y = center_index_y * 2.0 * self.semi_height + y_min

                    if point[1] > center_y:
                        x0, y0 = center_x + self.h_length, center_y
                        x1, y1 = center_x + semi_cell_length, center_y + self.semi_height
                        indicator = -(
                            point[1]
                            - (
                                (y0 - y1) / (x0 - x1) * point[0]
                                + (x0 * y1 - x1 * y0) / (x0 - x1)
                            )
                        )
                    else:
                        x0, y0 = center_x + semi_cell_length, center_y - self.semi_height
                        x1, y1 = center_x + self.h_length, center_y
                        indicator = point[1] - (
                            (y0 - y1) / (x0 - x1) * point[0]
                            + (x0 * y1 - x1 * y0) / (x0 - x1)
                        )

                    if indicator <= 0:
                        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 _hex_vertices(self, cx, cy):
            sh = self.semi_height
            lh = self.h_length
            return np.array(
                [
                    [cx + lh, cy],
                    [cx + lh / 2.0, cy + sh],
                    [cx - lh / 2.0, cy + sh],
                    [cx - lh, cy],
                    [cx - lh / 2.0, cy - sh],
                    [cx + lh / 2.0, cy - sh],
                    [cx + lh, cy],
                ],
                dtype=float,
            )

        def _iter_cells(self):
            dict_id_count = self.point_location_sta()

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

            for cell_id in dict_id_count.keys():
                ix = cell_id % self.count_column
                iy = cell_id // self.count_column

                cx = ix * self.h_length / 2.0 * 3.0 + x_min
                cy = iy * self.semi_height * 2.0 + y_min
                if ix % 2 == 1:
                    cy = (iy * 2.0 + 1.0) * self.semi_height + y_min

                verts = self._hex_vertices(cx, cy)
                poly = Polygon(verts[:-1])

                yield cell_id, poly, (cx, cy)

        def plot(
            self,
            title="Quadrat Count",
            show="counts",
            chi2_contrib=None,
        ):
            fig, ax = plt.subplots()
            ax.set_title(title)

            _scatter_points(ax, self.points)
            _draw_window(ax, self.window)

            dict_id_count = self.point_location_sta()

            patches = []
            ann_xy = []
            cell_ids = []

            for cell_id, poly, (cx, cy) in self._iter_cells():
                patches.append(MplPolygon(np.asarray(poly.exterior.coords), closed=True))
                ann_xy.append((cx, cy))
                cell_ids.append(cell_id)

            pc = PatchCollection(patches, linewidths=1.0, edgecolor="red", facecolor="none")
            ax.add_collection(pc)

            if show == "counts":
                for (cx, cy), cell_id in zip(ann_xy, cell_ids):
                    ax.text(
                        cx, cy, str(dict_id_count.get(cell_id, 0)), ha="center", va="center"
                    )

            elif show == "chi2":
                if chi2_contrib is None:
                    raise ValueError('chi2_contrib must be provided when show="chi2".')
                chi2_contrib = np.asarray(chi2_contrib, dtype=float)
                if chi2_contrib.shape[0] != len(cell_ids):
                    raise ValueError(
                        "chi2_contrib length must match number of plotted cells "
                        f"({len(cell_ids)}). Got {chi2_contrib.shape[0]}."
                    )
                pc.set_array(chi2_contrib)
                pc.set_alpha(0.6)
                pc.set_facecolor(None)
                plt.colorbar(pc, ax=ax, label="Chi-square contribution")

                for (cx, cy), v in zip(ann_xy, chi2_contrib):
                    ax.text(cx, cy, f"{v:.2f}", ha="center", va="center")

            else:
                raise ValueError('show must be "counts" or "chi2".')

            ax.set_aspect("equal", adjustable="box")
            return ax, cell_ids

    # -----------------------------------------------------------------------------
    # Quadrat statistic
    # -----------------------------------------------------------------------------

    class QStatistic:
        """Pearson chi-square quadrat test for complete spatial randomness (CSR).

        This class partitions the study region into quadrats (rectangles or hexagons),
        counts the number of events falling in each quadrat, and evaluates departure
        from CSR using a Pearson chi-square statistic under an equal-intensity CSR
        null model.

        The primary outputs are the observed chi-square statistic and its analytical
        p-value (from the chi-square distribution). Optionally, a simulation-based
        (Monte Carlo) p-value can be computed by generating CSR realizations within
        the study window and recomputing the chi-square statistic for each
        realization.

        Parameters
        ----------
        pp : :class: `.PointPattern` or array_like
            Event coordinates as an (n, 2) array-like of floats (x, y).
        shape : {"rectangle", "hexagon"}, default="rectangle"
            Quadrat tessellation type.
        nx : int, default=3
            Number of columns when `shape="rectangle"` and `rectangle_width` /
            `rectangle_height` are not provided.
        ny : int, default=3
            Number of rows when `shape="rectangle"` and `rectangle_width` /
            `rectangle_height` are not provided.
        rectangle_width : float, default=0
            Target rectangle width. Must be provided together with
            `rectangle_height`. When both are provided and non-zero, `nx` and `ny`
            are ignored and the grid dimensions are computed from the point-set MBB.
        rectangle_height : float, default=0
            Target rectangle height. Must be provided together with
            `rectangle_width`. When both are provided and non-zero, `nx` and `ny`
            are ignored and the grid dimensions are computed from the point-set MBB.
        lh : float, default=10
            Hexagon side length when `shape="hexagon"`.
        realizations : int, default=0
            Number of CSR simulations used to compute a Monte Carlo p-value.
            If 0, no simulation-based inference is performed.
        window : shapely geometry, optional
            Study window (e.g., Polygon or MultiPolygon). If None, the window is
            taken to be the axis-aligned MBB rectangle of `points`.
        rng : None | int | numpy.random.Generator | numpy.random.RandomState, optional
            Random number generator control for reproducible CSR simulations when
            `realizations > 0`.

            - None: create a new ``numpy.random.default_rng()``
            - int: use as a seed for ``numpy.random.default_rng(seed)``
            - Generator: used as-is
            - RandomState: wrapped via ``numpy.random.default_rng(RandomState)``

        Attributes
        ----------
        points : numpy.ndarray
            (n, 2) array of event coordinates.
        mbb : numpy.ndarray
            Bounding box of `points` as ``[xmin, ymin, xmax, ymax]``.
        window : shapely geometry
            Study window used for simulation and plotting.
        shape : str
            Quadrat tessellation type: "rectangle" or "hexagon".
        rng : numpy.random.Generator
            RNG used for CSR simulations (always stored as a Generator).
        mr : RectangleM or HexagonM
            Tessellation manager instance.
        cell_ids : list of int
            Cell identifiers included in the test statistic (all grid cells).
        chi2 : float
            Observed Pearson chi-square statistic.
        df : int
            Degrees of freedom for the analytical chi-square reference distribution,
            equal to ``k - 1`` where ``k`` is the number of included cells.
        chi2_pvalue : float
            Analytical p-value from the chi-square distribution.
        chi2_contrib : numpy.ndarray
            Per-cell chi-square contributions aligned with `cell_ids`, computed as
            ``(O - E)^2 / E`` with ``E = mean(O)`` over included cells.
        chi2_realizations : numpy.ndarray
            Simulated chi-square statistics for CSR realizations. Present only when
            `realizations > 0`.
        chi2_r_pvalue : float
            Monte Carlo p-value based on `chi2_realizations`, computed using the
            standard +1 correction:
            ``( #{T_sim >= T_obs} + 1 ) / (realizations + 1)``.
            Present only when `realizations > 0`.

        Notes
        -----
        - The analytical test uses ``scipy.stats.chisquare`` with expected
          counts equal across cells (i.e., ``E = mean(O)``). This
          corresponds to a homogeneous CSR null with equal-area cells. For
          irregular windows, a fully-correct analytical reference would
          area-weight expectations.

        - The simulation-based null generates CSR realizations within `window` with
          intensity ``n / area(window)``. Reproducibility depends on passing `rng`
          through to the underlying CSR generator (``poisson(..., rng=...)``).

        See Also
        --------
        RectangleM : Rectangular tessellation over the point-set MBB.
        HexagonM : Hexagonal tessellation over the point-set MBB.

        """

        def __init__(
            self,
            pp,
            shape="rectangle",
            nx=3,
            ny=3,
            rectangle_width=0,
            rectangle_height=0,
            lh=10,
            realizations=0,
            window=None,
            rng=None,
        ):
            self.points = _as_points_array(pp)
            self.mbb = _compute_mbb(self.points)
            self.window = _ensure_window(window, self.mbb)
            self.shape = shape
            self.rng = _coerce_rng(rng)

            if shape == "rectangle":
                self.mr = RectangleM(
                    self.points,
                    count_column=nx,
                    count_row=ny,
                    rectangle_width=rectangle_width,
                    rectangle_height=rectangle_height,
                    window=self.window,
                )
            elif shape == "hexagon":
                self.mr = HexagonM(self.points, lh, window=self.window)
            else:
                raise ValueError('shape must be either "rectangle" or "hexagon".')
            dict_id_count = self.mr.point_location_sta()

            obs_counts = np.asarray(list(dict_id_count.values()), dtype=float)
            self.cell_ids = list(dict_id_count.keys())

            self.chi2, self.chi2_pvalue = scipy.stats.chisquare(obs_counts)
            self.df = obs_counts.size - 1

            expected = obs_counts.mean() if obs_counts.size else np.nan
            if obs_counts.size and expected > 0:
                self.chi2_contrib = (obs_counts - expected) ** 2 / expected
            else:
                self.chi2_contrib = np.full_like(obs_counts, np.nan, dtype=float)

            # simulation-based inference under CSR (reproducible via rng)
            if realizations and realizations > 0:
                n = self.points.shape[0]
                intensity = n / float(self.window.area)

                reals = poisson(self.window, intensity, realizations, rng=self.rng)

                chi2_realizations = []
                for i in range(realizations):
                    ri = reals[i]
                    pts_i = getattr(ri, "points", ri)
                    pts_i = _as_points_array(pts_i)

                    if shape == "rectangle":
                        mr_temp = RectangleM(
                            pts_i,
                            count_column=self.mr.count_column,
                            count_row=self.mr.count_row,
                            window=self.window,
                        )
                    else:
                        mr_temp = HexagonM(pts_i, self.mr.h_length, window=self.window)

                    dtemp = mr_temp.point_location_sta()
                    sim_counts = np.asarray(list(dtemp.values()), dtype=float)

                    chi2_sim, _p = scipy.stats.chisquare(sim_counts)
                    chi2_realizations.append(chi2_sim)

                self.chi2_realizations = np.asarray(chi2_realizations, dtype=float)
                self.chi2_r_pvalue = (np.sum(self.chi2_realizations >= self.chi2) + 1.0) / (
                    realizations + 1.0
                )

        def plot(self, title="Quadrat Count", show="counts"):
            if show == "counts":
                ax, _cell_ids = self.mr.plot(title=title, show="counts")
                return ax

            if show == "chi2":
                # get plotted ids
                _, plotted_ids = self.mr.plot(title=title, show="counts")
                plt.close()

                contrib_map = {cid: v for cid, v in zip(self.cell_ids, self.chi2_contrib)}
                plotted_contrib = np.asarray(
                    [contrib_map.get(cid, np.nan) for cid in plotted_ids], dtype=float
                )

                ax, _ = self.mr.plot(title=title, show="chi2", chi2_contrib=plotted_contrib)
                return ax

            raise ValueError('show must be either "counts" or "chi2".')