# Source code for esda.lee

    import numpy
    from scipy import sparse
    from sklearn import preprocessing, utils
    from sklearn.base import BaseEstimator

    from .crand import _prepare_bivariate
    from .crand import njit as _njit

    class Spatial_Pearson(BaseEstimator):  # noqa: N801
        """Global Spatial Pearson Statistic"""

        def __init__(self, connectivity=None, permutations=999):
            """
            Initialize a spatial pearson estimator

            Parameters
            ----------
            connectivity:   scipy.sparse matrix object
                            the connectivity structure describing the relationships
                            between observed units. Will be row-standardized.
            permutations:   int
                            the number of permutations to conduct for inference.
                            if < 1, no permutational inference will be conducted.

            Attributes
            ----------
            association_: numpy.ndarray (2,2)
                          array containg the estimated Lee spatial pearson correlation
                          coefficients, where element [0,1] is the spatial correlation
                          coefficient, and elements [0,0] and [1,1] are the "spatial
                          smoothing factor"
            reference_distribution_: numpy.ndarray (n_permutations, 2,2)
                          distribution of correlation matrices for randomly-shuffled
                          maps.
            significance_: numpy.ndarray (2,2)
                           permutation-based p-values for the fraction of times the
                           observed correlation was more extreme than the simulated
                           correlations.
            """
            self.connectivity = connectivity
            self.permutations = permutations

        def fit(self, x, y):
            """
            bivariate spatial pearson's R based on Eq. 18 of :cite:`Lee2001`.

            L = \\dfrac{Z^T (V^TV) Z}{1^T (V^TV) 1}

            Parameters
            ----------
            x       :   numpy.ndarray
                        array containing continuous data
            y       :   numpy.ndarray
                        array containing continuous data

            Returns
            -------
            the fitted estimator.

            Notes
            -----
            Technical details and derivations can be found in :cite:`Lee2001`.

            """
            x = utils.check_array(x)
            y = utils.check_array(y)
            Z = numpy.column_stack(
                (
                    preprocessing.StandardScaler().fit_transform(x),
                    preprocessing.StandardScaler().fit_transform(y),
                )
            )
            if self.connectivity is None:
                self.connectivity = sparse.eye(Z.shape[0])
            self.association_ = self._statistic(Z, self.connectivity)

            if self.permutations is None or self.permutations < 1:
                return self

            if self.permutations:
                simulations = [
                    self._statistic(numpy.random.permutation(Z), self.connectivity)
                    for _ in range(self.permutations)
                ]
                self.reference_distribution_ = simulations = numpy.array(simulations)
                above = simulations >= self.association_
                larger = above.sum(axis=0)
                extreme = numpy.minimum(self.permutations - larger, larger)
                self.significance_ = (extreme + 1.0) / (self.permutations + 1.0)
            return self

        @staticmethod
        def _statistic(Z, W):
            ctc = W.T @ W
            ones = numpy.ones(ctc.shape[0])
            return (Z.T @ ctc @ Z) / (ones.T @ ctc @ ones)

    class Spatial_Pearson_Local(BaseEstimator):  # noqa: N801
        """Local Spatial Pearson Statistic"""

        def __init__(self, connectivity=None, permutations=999):
            """
            Initialize a spatial local pearson estimator

            Parameters
            ----------
            connectivity:   scipy.sparse matrix object
                            the connectivity structure describing the relationships
                            between observed units. Will be row-standardized.
            permutations:   int
                            the number of permutations to conduct for inference.
                            if < 1, no permutational inference will be conducted.
            significance_: numpy.ndarray (2,2)
                           permutation-based p-values for the fraction of times the
                           observed correlation was more extreme than the simulated
                           correlations.
            Attributes
            ----------
            associations_: numpy.ndarray (n_samples,)
                          array containg the estimated Lee spatial pearson correlation
                          coefficients, where element [0,1] is the spatial correlation
                          coefficient, and elements [0,0] and [1,1] are the "spatial
                          smoothing factor"
            reference_distribution_: numpy.ndarray (n_permutations, n_samples)
                          distribution of correlation matrices for randomly-shuffled
                          maps.
            significance_: numpy.ndarray (n_samples,)
                           permutation-based p-values for the fraction of times the
                           observed correlation was more extreme than the simulated
                           correlations.

            Notes
            -----
            Technical details and derivations can be found in :cite:`Lee2001`.

            """
            self.connectivity = connectivity
            self.permutations = permutations

        def fit(self, x, y):
            """
            Bivariate local pearson's R based on Eq. 22 in Lee (2001), using
            site-wise conditional randomization from Moran_Local_BV.

            .. math::

                L_i = \\dfrac{
                             n \\cdot
                               \\Big[\big(\\sum_i w_{ij}(x_j - \bar{x})\big)
                                     \big(\\sum_i w_{ij}(y_j - \bar{y})\big) \\Big]
                             }
                            {
                             \\sqrt{\\sum_i (x_i - \bar{x})^2}
                             \\sqrt{\\sum_i (y_i - \bar{y})^2}}
                    = \\dfrac{
                             n \\cdot
                               (\tilde{x}_j - \bar{x})
                               (\tilde{y}_j - \bar{y})
                             }
                            {
                             \\sqrt{\\sum_i (x_i - \bar{x})^2}
                             \\sqrt{\\sum_i (y_i - \bar{y})^2}}

            Lee, Sang Il. (2001), "Developing a bivariate spatial
            association measure: An integration of Pearson's r and
            Moran's I." Journal of Geographical Systems, 3(4):369-385.

            Parameters
            ----------
            x       :   numpy.ndarray
                        array containing continuous data
            y       :   numpy.ndarray
                        array containing continuous data

            Returns
            -------
            the fitted estimator.
            """
            x = utils.check_array(x)
            x = preprocessing.StandardScaler().fit_transform(x)

            y = utils.check_array(y)
            y = preprocessing.StandardScaler().fit_transform(y)

            Z = numpy.column_stack((x, y))

            standard_connectivity = sparse.csc_matrix(
                self.connectivity / self.connectivity.sum(axis=1).reshape(-1, 1)
            )

            n, _ = x.shape

            self.associations_ = self._statistic(Z, standard_connectivity)

            if self.permutations:
                self.reference_distribution_ = numpy.empty((n, self.permutations))
                max_neighbors = (standard_connectivity != 0).sum(axis=1).max()
                random_ids = numpy.array(
                    [
                        numpy.random.permutation(n - 1)[0 : max_neighbors + 1]
                        for i in range(self.permutations)
                    ]
                )
                ids = numpy.arange(n)

                for i in range(n):
                    row = standard_connectivity[i]
                    weight = numpy.asarray(row[row.nonzero()]).reshape(-1, 1)
                    cardinality = row.nonzero()[0].shape[0]

                    ids_not_i = ids[ids != i]
                    numpy.random.shuffle(ids_not_i)
                    randomizer = random_ids[:, 0:cardinality]
                    random_neighbors = ids_not_i[randomizer]

                    random_neighbor_x = x[random_neighbors]
                    random_neighbor_y = y[random_neighbors]

                    self.reference_distribution_[i] = (
                        (weight * random_neighbor_y).sum(axis=1) - y.mean()
                    ).squeeze()
                    self.reference_distribution_[i] *= (
                        (weight * random_neighbor_x).sum(axis=1) - x.mean()
                    ).squeeze()
                above = self.reference_distribution_ >= self.associations_.reshape(-1, 1)
                larger = above.sum(axis=1)
                extreme = numpy.minimum(larger, self.permutations - larger)
                self.significance_ = (extreme + 1.0) / (self.permutations + 1.0)
                self.reference_distribution_ = self.reference_distribution_.T
            else:
                self.reference_distribution_ = None
            return self

        @staticmethod
        def _statistic(Z, W):
            return (Z[:, 1] @ W.T) * (W @ Z[:, 0])

    # --------------------------------------------------------------
    # Conditional Randomization Function Implementations
    # --------------------------------------------------------------

    @_njit(fastmath=True)
    def _local_spatial_pearson_crand(i, z, permuted_ids, weights_i, scaling):
        zxi, zxrand, zyi, zyrand = _prepare_bivariate(i, z, permuted_ids, weights_i)
        return (zyrand @ weights_i) * (zxrand @ weights_i) * scaling

    if __name__ == "__main__":
        import geopandas
        import libpysal

        df = geopandas.read_file(libpysal.examples.get_path("columbus.shp"))
        x = df[["HOVAL"]].values
        y = df[["CRIME"]].values
        zx = preprocessing.StandardScaler().fit_transform(x)
        zy = preprocessing.StandardScaler().fit_transform(y)
        w = libpysal.weights.Queen.from_dataframe(df)
        w.transform = "r"
        numpy.random.seed(2478879)
        testglobal = Spatial_Pearson(connectivity=w.sparse).fit(x, y)
        numpy.random.seed(2478879)
        testlocal = Spatial_Pearson_Local(connectivity=w.sparse).fit(x, y)