
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import jax.numpy as jnp
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import scipy.sparse


# Import from core
from .core import constants
from .core.constants import TOLERANCE, EPSILON, GEOMETRY_EPSILON, PLOT_LOG_BIAS


def dist_sqeuclidean(XA, XB):
    """JAX-based squared Euclidean pairwise distance calculation.
    
    Args:
        XA: array of shape (m, n)
        XB: array of shape (p, n)
    
    Returns:
        Distance matrix of shape (m, p)
    """
    XA = jnp.asarray(XA, dtype=constants.DTYPE_FLOAT)
    XB = jnp.asarray(XB, dtype=constants.DTYPE_FLOAT)
    # Squared Euclidean: ||a-b||^2 = ||a||^2 + ||b||^2 - 2*a·b
    XA_sq = jnp.sum(XA**2, axis=1, keepdims=True)
    XB_sq = jnp.sum(XB**2, axis=1, keepdims=True)
    return XA_sq + XB_sq.T - 2 * jnp.dot(XA, XB.T)

def consistent_cos(thetas_deg):
    """
    consistent cosine function for angles in degrees
    reduces angle to first quadrant of unit circle

    The need for this function is that the jax cosine function does not have
    cos(90deg) = exactly 0.0, or cos(x)+cos(180deg+x) = exactly 0.0
    
    args:
        thetas_deg: array of angles in degrees

    returns:
        array of cosines
    """

    thetas_deg = jnp.abs(thetas_deg)
    thetas_deg = thetas_deg % 360
    signs_180 = jnp.where(thetas_deg>=180.0, -1.0, 1.0)
    thetas_deg = jnp.where(thetas_deg>=180.0, thetas_deg-180.0, thetas_deg)
    signs_90 = jnp.where(thetas_deg>90.0, -1.0, 1.0)
    thetas_deg = jnp.where(thetas_deg>90.0, 180.0-thetas_deg, thetas_deg)
    zeros_90 = jnp.where(thetas_deg==90.0, 0.0, 1.0)
    return zeros_90 * signs_180 * signs_90 * jnp.cos(jnp.deg2rad(thetas_deg)) 

def consistent_sin(thetas_deg):
    """
    consistent sine function for angles in degrees
    reduces angle to first quadrant of unit circle

    The need for this function is that the jax sin function does not have
    sin(90deg) = exactly 1.0, sin(x)+sin(180deg+x) = exactly 0.0
    
    args:
        thetas_deg: array of angles in degrees

    returns:
        array of cosines
    """

    signs_neg = jnp.where(thetas_deg<0.0, -1.0, 1.0)
    thetas_deg = jnp.where(thetas_deg<0.0, -thetas_deg, thetas_deg)
    thetas_deg = thetas_deg % 360
    signs_180 = jnp.where(thetas_deg>=180.0, -1.0, 1.0)
    thetas_deg = jnp.where(thetas_deg>=180.0, thetas_deg-180.0, thetas_deg)
    thetas_deg = jnp.where(thetas_deg>90.0, 180.0-thetas_deg, thetas_deg)
    return signs_neg * signs_180 * jnp.where(thetas_deg==90.0, 1.0, jnp.sin(jnp.deg2rad(thetas_deg))) 

def polar_to_cartesian(r_theta_deg):
    """Converts polar coordinates to Cartesian coordinates.
    
    Args:
        r_theta_deg: array of shape (n, 2), where each row is (r, theta_deg)
    
    Returns:
        array of shape (n, 2), where each row is (x, y)
    """
    r_theta_deg = jnp.asarray(r_theta_deg, dtype=constants.DTYPE_FLOAT)
    r, theta_deg = r_theta_deg[:,0], r_theta_deg[:,1]
    x = r * consistent_cos(theta_deg)
    y = r * consistent_sin(theta_deg)
    return jnp.column_stack((x,y))

def dist_sqeuclidean_polar(r_theta_A, r_theta_B):
    """JAX-based squared Euclidean pairwise distance calculation in polar coordinates.
    
    Args:
        r_theta_A: array of shape (m, 2), where each row is (r, theta_deg)
        r_theta_B: array of shape (p, 2), where each row is (r, theta_deg)
    
    Returns:
        Distance matrix of shape (m, p)
    """
    r_theta_A = jnp.asarray(r_theta_A, dtype=constants.DTYPE_FLOAT)
    r_theta_B = jnp.asarray(r_theta_B, dtype=constants.DTYPE_FLOAT)
    r_A, theta_A = r_theta_A[:,0], r_theta_A[:,1]
    r_B, theta_B = r_theta_B[:,0], r_theta_B[:,1]
    # Squared Euclidean: ||a-b||^2 = ||a||^2 + ||b||^2 - 2*a·b
    r_A_sq = r_A**2
    r_B_sq = r_B**2
    return r_A_sq[:, None] + r_B_sq[None, :] - 2 * r_A[:, None] * r_B[None, :] * consistent_cos(theta_A[:, None] - theta_B[None, :])

def dist_manhattan(XA, XB):
    """JAX-based Manhattan pairwise distance calculation.
    
    Args:
        XA: array of shape (m, n)
        XB: array of shape (p, n)
    
    Returns:
        Distance matrix of shape (m, p)
    """
    XA = jnp.asarray(XA, dtype=constants.DTYPE_FLOAT)
    XB = jnp.asarray(XB, dtype=constants.DTYPE_FLOAT)
    # Manhattan distance: sum(|a-b|)
    return jnp.sum(jnp.abs(XA[:, None, :] - XB[None, :, :]), axis=2)


def _is_in_triangle_single(p, a, b, c):
    """
    Returns True if point p is in triangle (a, b, c).
    Robust for arbitrary vertex winding (CW or CCW).
    
    Args:
        p: Point as [x, y]
        a, b, c: Triangle vertices as [x, y]
    
    Returns:
        Boolean indicating if p is inside triangle

    See also:  computational geometry, half-plane test;
    Stack Overflow answer to https://stackoverflow.com/questions/2049582/how-to-determine-if-a-point-is-in-a-2d-triangle
       https://stackoverflow.com/a/2049593/103081 
       by https://stackoverflow.com/users/233522/kornel-kisielewicz
    """
    def cross(o, a, b):
        return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])

    s1 = cross(p, a, b)
    s2 = cross(p, b, c)
    s3 = cross(p, c, a)

    # Use centralized epsilon from core
    eps = GEOMETRY_EPSILON
    has_neg = (s1 < -eps) | (s2 < -eps) | (s3 < -eps)
    has_pos = (s1 > eps) | (s2 > eps) | (s3 > eps)
    
    return ~(has_neg & has_pos)


class Grid:
    def __init__(self, *, x0, x1, xstep=1, y0, y1, ystep=1):
        """inititalizes a rectangular Grid object 
        
        Args:
            x0: minimum x value
            x1: maximum x value
            xstep: step size for x
            y0: minimum y value
            y1: maximum y value
            ystep: step size for y
        

        Properties:
            x0: minimum x value
            x1: maximum x value
            xstep: step size for x
            y0: minimum y value
            y1: maximum y value
            ystep: step size for y
            extent: (x0, x1, y0, y1)
            gshape: (number_of_rows, number_of_cols)
            len: number_of_rows * number_of_cols
            x: 1D JAX array of x coordinates
            y: 1D JAX array of y coordinates
            points: 2D JAX array of (x,y) coordinates
            boundary: 1D JAX boolean array for boundary points
            weights: 1D JAX array of weights
        
        Returns:
            None
        """
        self.x0 = x0
        self.y0 = y0
        self.x1 = x1
        self.y1 = y1
        self.xstep = xstep
        self.ystep = ystep
        self.extent = (self.x0, self.x1, self.y0, self.y1)
        self.gshape = self.shape()
        self.len = self.gshape[0] * self.gshape[1]
        def _coords(i):
            """returns the (x,y) coordinate tuple of the i-th grid point
            
            Args:
                i: index of the grid point
            Returns:
                (x,y) coordinate tuple of the i-th grid point

            Note: this is a helper function for __init__, and is not defined as a method to avoid polluting subclasses with this function    
            """
            rows, cols = self.gshape
            row = i // cols
            col = i % cols
            return (self.x0 + col * self.xstep, self.y1 - row * self.ystep)
        self.x, self.y = _coords(jnp.arange(self.len))
        self.points = jnp.column_stack((self.x,self.y))
        self.boundary = ((self.x==x0) | (self.x==x1) | (self.y==y0) | (self.y==y1))
        self.weights = None

    def shape(self, *, x0=None, x1=None, xstep=None, y0=None, y1=None, ystep=None):
        """returns a tuple(number_of_rows,number_of_cols) for the natural shape of the current grid, or a subset"""
        x0 = self.x0 if x0 is None else x0
        x1 = self.x1 if x1 is None else x1
        y0 = self.y0 if y0 is None else y0
        y1 = self.y1 if y1 is None else y1
        xstep = self.xstep if xstep is None else xstep
        ystep = self.ystep if ystep is None else ystep
        if x1 < x0:
            raise ValueError
        if y1 < y0:
            raise ValueError
        if xstep <= 0:
            raise ValueError
        if ystep <= 0:
            raise ValueError
        number_of_rows = 1 + int((y1 - y0) / ystep)
        number_of_cols = 1 + int((x1 - x0) / xstep)
        return (number_of_rows, number_of_cols)

    def within_box(self, *, x0=None, x1=None, y0=None, y1=None):
        """returns a 1D numpy boolean array, suitable as an index mask, for testing whether a grid point is also in the defined box"""
        x0 = self.x0 if x0 is None else x0
        x1 = self.x1 if x1 is None else x1
        y0 = self.y0 if y0 is None else y0
        y1 = self.y1 if y1 is None else y1
        return (self.x >= x0) & (self.x <= x1) & (self.y >= y0) & (self.y <= y1)

    def within_disk(self, *, x0, y0, r, metric="euclidean", **kwargs):
        """returns 1D JAX boolean array, suitable as an index mask, for testing whether a grid point is also in the defined disk"""
        center = jnp.array([[x0, y0]])
        
        if metric == "euclidean":
            # For Euclidean distance, use squared Euclidean and compare r^2
            distances_sq = dist_sqeuclidean(center, self.points)
            mask = (distances_sq <= r**2).flatten()
        elif metric == "manhattan":
            distances = dist_manhattan(center, self.points)
            mask = (distances <= r).flatten()
        else:
            raise ValueError(f"Unsupported metric: {metric}. Use 'euclidean' or 'manhattan'.")
        
        return mask
    
    def within_triangle(self, *, points):
        """returns 1D JAX boolean array, suitable as an index mask, for testing whether a grid point is also in the defined triangle"""
        points = jnp.asarray(points)
        a, b, c = points[0], points[1], points[2]
        
        # Vectorized cross-product triangle containment test
        # Use vmap to apply the single-point test to all grid points
        mask = jax.vmap(
            lambda p: _is_in_triangle_single(p, a, b, c)
        )(self.points)
        
        return mask

    def index(self, *, x, y, tolerance=1e-9):
        """
        Returns the unique 1D array index for grid point (x,y).
        
        Uses direct computation for O(1) lookup instead of linear search.
        For regular grid: index = row * n_cols + col
        where row = (y1 - y) / ystep, col = (x - x0) / xstep
        
        Args:
            x: x-coordinate
            y: y-coordinate
            tolerance: tolerance for coordinate matching (default: 1e-9)
        
        Returns:
            int: Grid index, or raises ValueError if point not on grid
        """
        # Compute row and column indices
        col = jnp.round((x - self.x0) / self.xstep)
        row = jnp.round((self.y1 - y) / self.ystep)
        
        # Check if within bounds
        n_rows, n_cols = self.gshape
        if not (0 <= row < n_rows and 0 <= col < n_cols):
            raise ValueError(f"Point ({x}, {y}) is outside grid bounds")
        
        # Compute index
        idx = row * n_cols + col
        
        # Verify the point matches (within tolerance)
        if abs(self.x[idx] - x) > tolerance or abs(self.y[idx] - y) > tolerance:
            raise ValueError(f"Point ({x}, {y}) does not match grid point at computed index")
        
        return int(idx)

    def embedding(self, *, valid):
        """
        returns an embedding function efunc(z,fill=0.0) from 1D arrays z of size sum(valid)
        to arrays of size self.len

        valid is a jnp.array of type boolean, of size self.len

        fill is the value for indices outside the embedding. The default
        is zero (0.0).  Setting fill=jnp.nan can be useful for
        plotting purposes as matplotlib will omit jnp.nan values from various
        kinds of plots.
        """

        correct_z_len = valid.sum()

        def efunc(z, fill=0.0):
            v = jnp.full(self.len, fill)
            return v.at[valid].set(z)

        return efunc

    def extremes(self, z, *, valid=None):
        """
        Returns the minimum and maximum values of z, along with point arrays for argmin and argmax
        
        Args:
            z: Array of values of size self.len or valid.len
            valid: Boolean array indicating which grid points are valid
        
        Returns: tuple (min_z, min_z_points, max_z, max_z_points)
            min_z: Minimum value of z
            min_z_points: Array of points where z is equal to min_z
            max_z: Maximum value of z
            max_z_points: Array of points where z is equal to max_z

        Note: The relative tolerance for min and max is 2*EPSILON, where 
        EPSILON is the machine epsilon set for float32 or float64 in core/constants.py

        Note: uses dynamic indexing, so cannot be used in jax.jit
        """
        # if valid is None return unrestricted min,points_min,max,points_max
        # if valid is a boolean array, return constrained min,points_min,max,points_max
        # note that min/max is always calculated over all of z, it is the points that must be restricted
        # because valid indicates that z came from a subset of the points
        min_z = float(z.min())
        min_z_mask = jnp.abs(z-min_z) <= 2*EPSILON*jnp.abs(min_z)
        max_z = float(z.max())
        max_z_mask = jnp.abs(z-max_z) <= 2*EPSILON*jnp.abs(max_z)
        if valid is None:
           return (min_z,self.points[min_z_mask],max_z,self.points[max_z_mask]) 
        return (min_z,self.points[valid][min_z_mask],max_z,self.points[valid][max_z_mask])

    def spatial_utilities(
        self, *, voter_ideal_points, metric="sqeuclidean", scale=-1, decimals=None
    ):
        """
        returns utility function values for each voter at each grid point
        
        Args:
            voter_ideal_points: Array of voter ideal points of shape (n_voters, 2)
            metric: Metric to use for distance calculation (default: "sqeuclidean" or "manhattan")
            scale: Scale factor for distance (default: -1)
            decimals: Number of decimals to round utility functions to (default: no rounding)
        """
        voter_ideal_points = jnp.asarray(voter_ideal_points, dtype=constants.DTYPE_FLOAT)
        
        if metric == "sqeuclidean":
            distances = dist_sqeuclidean(voter_ideal_points, self.points)
        elif metric == "manhattan":
            distances = dist_manhattan(voter_ideal_points, self.points)
        else:
            raise ValueError(f"Unsupported metric: {metric}. Use 'sqeuclidean' or 'manhattan'.")
        
        if decimals is not None:
            return jnp.round(scale * distances, decimals=decimals)
        return scale * distances

    def plot(
        self,
        z,
        *,
        title=None,
        cmap=cm.gray_r,
        alpha=0.6,
        alpha_points=0.3,
        log=True,
        points=None,
        zoom=False,
        border=1,
        logbias=PLOT_LOG_BIAS, # Use constant from core
        figsize=(10, 10),
        dpi=72,
        fname=None
    ):
        """plots values z defined on the grid;
        optionally plots additional 2D points
         and zooms to fit the bounding box of the points"""
        # Convert JAX arrays to NumPy for matplotlib compatibility
        z = np.array(z)
        grid_x = np.array(self.x)
        grid_y = np.array(self.y)
        
        plt.figure(figsize=figsize, dpi=dpi)
        plt.rcParams["font.size"] = "24"
        fmt = "%1.2f" if log else "%.2e"
        if zoom:
            points = np.asarray(points)
            [min_x, min_y] = np.min(points, axis=0) - border
            [max_x, max_y] = np.max(points, axis=0) + border
            box = {"x0": min_x, "x1": max_x, "y0": min_y, "y1": max_y}
            inZoom = np.array(self.within_box(**box))
            zshape = self.shape(**box)
            extent = (min_x, max_x, min_y, max_y)
            zraw = np.copy(z[inZoom]).reshape(zshape)
            x = np.copy(grid_x[inZoom]).reshape(zshape)
            y = np.copy(grid_y[inZoom]).reshape(zshape)
        else:
            zshape = self.gshape
            extent = self.extent
            zraw = z.reshape(zshape)
            x = grid_x.reshape(zshape)
            y = grid_y.reshape(zshape)
        zplot = np.log10(logbias + zraw) if log else zraw
        contours = plt.contour(x, y, zplot, extent=extent, cmap=cmap)
        plt.clabel(contours, inline=True, fontsize=12, fmt=fmt)
        plt.imshow(zplot, extent=extent, cmap=cmap, alpha=alpha)
        if points is not None:
            plt.scatter(points[:, 0], points[:, 1], alpha=alpha_points, color="black")
        if title is not None:
            plt.title(title)
        if fname is None:
            plt.show()
        else:
            plt.savefig(fname)


    def parts_from_linear_discriminator(self, center=(0,0), d1=(0,0), d2=(0,0), d3=(0,0)) -> jnp.ndarray:
        """
        Generate partition using linear discriminant function.
    
        Creates partition by evaluating a discriminant function at each grid point:
        f(p) = (p-center)·d1 + |(p-center)|·d2 + |(p-center)·d3|
    
        Points with equal discriminant values are grouped together.
    
        Args:
            center: Center point for discriminant (default: (0,0))
            d1: Linear term direction vector (default: (0,0))
            d2: Absolute value term direction vector (default: (0,0))
            d3: Absolute dot product direction vector (default: (0,0))
        
        Returns:
            Inverse indices array grouping points by discriminant value
        
        Notes:
            - Used internally for efficient symmetry partitioning
            - Optimized for reflection symmetries on regular grids
        """
        center = jnp.array(center)
        d1 = jnp.array(d1)
        d2 = jnp.array(d2)
        d3 = jnp.array(d3)
        centered = self.points - center
        discriminant_values = jnp.dot(centered, d1) + jnp.dot(jnp.abs(centered), d2) + jnp.abs(jnp.dot(centered, d3))
        inverse_indices = jnp.unique(discriminant_values, return_inverse=True)[1]
        return inverse_indices

    def partition_from_symmetry(
        self,
        symmetries: list,
        tolerance: float = 1e-6
    ) -> jnp.ndarray:
        """
        Generate partition from spatial symmetries.
        
        Builds partition by grouping grid points that are equivalent under
        the specified spatial symmetries. Does not verify symmetry in the
        transition matrix - assumes user-specified symmetries are correct.
        
        Args:
            symmetries: List of symmetry specifications:
                - 'reflect_x' or 'reflect_x=0': Reflection around x=0
                - 'reflect_x=c': Reflection around x=c
                - 'reflect_y' or 'reflect_y=0': Reflection around y=0
                - 'reflect_y=c': Reflection around y=c
                - 'reflect_xy': Reflection around line y=x
                - 'swap_xy': Swap x and y coordinates (equivalent to reflect_xy)
                - ('rotate', center_x, center_y, degrees): Rotation around (cx, cy)
                  Example: ('rotate', 0, 0, 120) for 120° rotation around origin
            tolerance: Distance tolerance for matching rotated points (default: 1e-6)
                       Useful for approximate symmetries like 120° rotation on grid
        
        Returns:
            jnp.ndarray: Inverse indices array where inverse_indices[i] gives
                        the group ID for grid point i
        
        Examples:
            >>> # Reflection symmetry around y-axis
            >>> partition = grid.partition_from_symmetry(['reflect_x'])
            
            >>> # (x,y) <-> (y,x) symmetry
            >>> partition = grid.partition_from_symmetry(['swap_xy'])
            
            >>> # 120° rotation (BJM spatial triangle example)
            >>> # Grid points near 120° rotations are grouped
            >>> partition = grid.partition_from_symmetry(
            ...     [('rotate', 0, 0, 120)], tolerance=0.5
            ... )
        
        Notes:
            - Symmetries are applied iteratively to build equivalence classes
            - Does not validate that the Markov chain respects these symmetries
            - Rotation tolerance allows approximate symmetries
            - User is responsible for ensuring symmetries are appropriate
            - Optimized for regular grids using direct index computation
            - Use suggest_symmetries() from gridvoting_jax.symmetry to automatically
              detect symmetries in voter ideal points
        """
        n_states = self.len
        
        # Fast path: For singleton symmetries, use parts_from_linear_discriminator
        if len(symmetries) == 1:
            sym = symmetries[0]
            
            # Handle string symmetries
            if isinstance(sym, str):
                if sym == 'reflect_x' or sym.startswith('reflect_x='):
                    # Extract offset if present
                    offset = 0.0 if sym == 'reflect_x' else float(sym.split('=')[1])
                    # Reflection across vertical line x=offset: (x,y) ≡ (2*offset-x, y)
                    # Discriminant: self.len * y + |x - offset|
                    return self.parts_from_linear_discriminator(
                        center=(offset, 0),
                        d1=(0, self.len),  # self.len * y
                        d2=(1, 0),         # |x - offset|
                        d3=(0, 0)
                    )
                
                elif sym == 'reflect_y' or sym.startswith('reflect_y='):
                    # Extract offset if present
                    offset = 0.0 if sym == 'reflect_y' else float(sym.split('=')[1])
                    # Reflection across horizontal line y=offset: (x,y) ≡ (x, 2*offset-y)
                    # Discriminant: self.len * x + |y - offset|
                    return self.parts_from_linear_discriminator(
                        center=(0, offset),
                        d1=(self.len, 0),  # self.len * x
                        d2=(0, 1),         # |y - offset|
                        d3=(0, 0)
                    )
                
                elif sym in ('swap_xy', 'reflect_xy'):
                    # Diagonal reflection: (x,y) ≡ (y,x)
                    # Discriminant: self.len * (x+y) + |x-y|
                    return self.parts_from_linear_discriminator(
                        center=(0, 0),
                        d1=(self.len, self.len),  # self.len * (x+y)
                        d2=(0, 0),
                        d3=(1, -1)                # |x-y|
                    )
        
        # General case: Use connected components for multiple symmetries or rotations
        # We will build a list of edges (u, v) representing symmetric equivalence
        # source_indices = []
        # target_indices = []
        
        # Use JAX/NumPy for vectorized coordinate transformation
        # x, y are standard numpy arrays (or JAX arrays)
        X = self.x
        Y = self.y
        
        # Accumulate edges in efficient list of arrays
        edges_src = []
        edges_dst = []
        
        # Always include self-loops to ensure every node is in the graph
        # (though connected_components handles isolated nodes, explicitly adding identity is safe)
        edges_src.append(jnp.arange(n_states))
        edges_dst.append(jnp.arange(n_states))

        for sym in symmetries:
            # 1. Vectorized Transformation
            # ----------------------------------------------------------------
            if isinstance(sym, str):
                if sym == 'swap_xy' or sym == 'reflect_xy':
                    # Swap: (x, y) -> (y, x)
                    X_new, Y_new = Y, X
                
                elif sym.startswith('reflect_x'):
                    # Reflect x around c: x' = 2c - x
                    c = float(sym.split('=')[1]) if '=' in sym else 0.0
                    X_new = 2 * c - X
                    Y_new = Y
                
                elif sym.startswith('reflect_y'):
                    # Reflect y around c: y' = 2c - y
                    c = float(sym.split('=')[1]) if '=' in sym else 0.0
                    X_new = X
                    Y_new = 2 * c - Y
                else:
                    raise ValueError(f"Unknown symmetry string: {sym}")
            
            elif isinstance(sym, tuple) and sym[0] == 'rotate':
                # Rotate around (cx, cy) by degrees
                _, cx, cy, degrees = sym
                theta = np.radians(degrees)
                cos_t = np.cos(theta)
                sin_t = np.sin(theta)
                
                # Apply rotation matrix
                dx = X - cx
                dy = Y - cy
                X_new = cx + (dx * cos_t - dy * sin_t)
                Y_new = cy + (dx * sin_t + dy * cos_t)
            else:
                 raise ValueError(f"Unknown symmetry spec: {sym}")

            # 2. Vectorized Index Lookup (Regular Grid)
            # ----------------------------------------------------------------
            # Expected indices (nearest integer grid point)
            # col = (x - x0) / xstep
            # row = (y1 - y) / ystep  (note y1 is top)
            
            # Use jnp.rint (round to nearest integer)
            col_new = jnp.rint((X_new - self.x0) / self.xstep).astype(jnp.int32)
            row_new = jnp.rint((self.y1 - Y_new) / self.ystep).astype(jnp.int32)
            
            # 3. Filtering
            # ----------------------------------------------------------------
            n_rows, n_cols = self.gshape
            
            # Check bounds
            mask_bounds = (col_new >= 0) & (col_new < n_cols) & \
                          (row_new >= 0) & (row_new < n_rows)
            
            # Calculate 1D index for potentially valid points
            # We must be careful not to access invalid indices, so we apply mask immediately
            # But JAX supports careful masking.
            # Let's compute hypothetical indices, then filter.
            idx_new = row_new * n_cols + col_new
            
            # Check coordinate match (tolerance)
            # Only check where bounds are valid to avoid OOB indexing
            # For OOB, we set error distance to infinity or just mask them out first
            
            # Strategy: Filter indices first
            valid_indices = jnp.where(mask_bounds)[0]
            target_indices = idx_new[valid_indices]
            
            # Check distance on valid candidates
            # X_target = self.x[target_indices]
            # Y_target = self.y[target_indices]
            dist_x = jnp.abs(self.x[target_indices] - X_new[valid_indices])
            dist_y = jnp.abs(self.y[target_indices] - Y_new[valid_indices])
            
            mask_match = (dist_x <= tolerance) & (dist_y <= tolerance)
            
            # Final valid edges
            # Source: valid_indices[mask_match]
            # Target: target_indices[mask_match]
            
            final_src = valid_indices[mask_match]
            final_dst = target_indices[mask_match]
            
            edges_src.append(final_src)
            edges_dst.append(final_dst)
            
        # 4. Graph Construction & Partitioning (CPU/SciPy)
        # ----------------------------------------------------------------
        # Concatenate all edges
        all_src = jnp.concatenate(edges_src)
        all_dst = jnp.concatenate(edges_dst)
        
        # Convert to numpy for SciPy
        all_src_np = np.array(all_src)
        all_dst_np = np.array(all_dst)
        
        # Build Sparse Matrix (Adjacency)
        # We need a symmetric graph for connected components, but csgraph handles directed/undirected
        # connection_type='weak' treats directed edges as undirected for components
        # Weights don't matter, just connectivity. Use 1s.
        data = np.ones(len(all_src_np), dtype=bool)
        
        adj = scipy.sparse.coo_matrix(
            (data, (all_src_np, all_dst_np)),
            shape=(n_states, n_states)
        )
        
        # Find Connected Components
        # connection='weak' means if u->v or v->u, they are connected.
        # This is correct because symmetry implies equivalence A~B.
        n_components, labels = scipy.sparse.csgraph.connected_components(
            adj, 
            directed=True, 
            connection='weak',
            return_labels=True
        )
        
        # Convert labels to inverse indices (already in correct format!)
        inverse_indices = jnp.array(labels, dtype=jnp.int32)
        
        return inverse_indices


class PolarGrid(Grid):
    def __init__(self, *, radius, rstep=1, thetastep=15):
        """initializes PolarGrid object
        
        Args:
            radius: radius of the grid
            rstep: step size for the radial coordinate
            thetastep: step size for the angular coordinate

        Properties:
            radius: radius of the grid
            rstep: step size for the radial coordinate
            thetastep: step size for the angular coordinate (degrees)
            rvals: radial coordinate values
            thetavals: angular coordinate values (degrees)
            n_rvals: number of radial coordinate values
            n_thetavals: number of angular coordinate values
            len: total number of grid points
            r: radial coordinates for each grid point
            theta_deg: angular coordinates (degrees) for each grid point
            r_theta_deg: radial and angular coordinates (degrees) for each grid point
            boundary: boolean array for boundary points
            weights: planar area around each grid point
            x0,y0,x1,y1: rectangular extent of the grid
        """
        self.radius = radius
        self.rstep = rstep
        self.thetastep = thetastep
        self.rvals = jnp.arange(0, radius + rstep, rstep)
        self.thetavals = jnp.arange(0, 360, thetastep)
        self.n_rvals = len(self.rvals)
        self.n_thetavals = len(self.thetavals)
        self.len = 1 + ((self.n_rvals-1) * self.n_thetavals)
        self.r = jnp.concat((jnp.array([0.0]),jnp.repeat(self.rvals[1:], self.n_thetavals)))
        assert self.r.shape[0] == self.len 
        self.theta_deg = jnp.concat((jnp.array([0.0]),jnp.tile(self.thetavals, self.n_rvals-1)))
        assert self.theta_deg.shape[0] == self.len
        self.r_theta_deg = jnp.column_stack((self.r, self.theta_deg))
        self.points = polar_to_cartesian(self.r_theta_deg)
        self.boundary = self.r==self.radius
        # weights are the area of each grid cell
        self.weights = (
            (jnp.square(jnp.where(self.r+0.5*self.rstep>self.radius, self.radius, self.r+0.5*self.rstep))
            -jnp.square(self.r-0.5*self.rstep)
            )*jnp.deg2rad(self.thetastep)/2.0
            ).at[0].set(0.25*jnp.pi*self.rstep**2)
        total_weight = self.weights.sum()
        expected_weight = jnp.pi*self.radius**2
        assert jnp.allclose(total_weight, expected_weight, rtol=1e-4), f"Total weight {total_weight} does not match expected weight {expected_weight}"
        self.x0 = -self.radius
        self.x1 = self.radius
        self.y0 = -self.radius
        self.y1 = self.radius
        self.extent = (self.x0, self.x1, self.y0, self.y1)

    def shape(self):
        """
        Not implemented for PolarGrid
        """
        raise NotImplementedError
    
    def spatial_utilities(
        self, *, voter_ideal_points, metric="sqeuclidean", scale=-1, decimals=None
    ):
        """
        returns utility function values for each voter at each grid point
        
        Args:
            voter_ideal_points: Array of voter ideal points of shape (n_voters, 2)
            metric: Metric to use for distance calculation (default: "sqeuclidean")
            scale: Scale factor for distance (default: -1)
            decimals: Number of decimals to round utility functions to (default: no rounding)
        """
        voter_ideal_points = jnp.asarray(voter_ideal_points, dtype=constants.DTYPE_FLOAT)
        
        if metric == "sqeuclidean":
            distances = dist_sqeuclidean_polar(voter_ideal_points, self.r_theta_deg)
        else:
            raise ValueError(f"Unsupported metric: {metric}. Use 'sqeuclidean'")
        
        if decimals is not None:
            return jnp.round(scale * distances, decimals=decimals)
        return scale * distances

    def as_rings(self, *, z):
        """
        Returns the z values reshaped into rings of constant radius
        Args:
            z: z values, a jax array of length self.len
        Returns:
            tuple of (z[0], 2D matrix of z values with rows indicating r and columns indicating theta)

        Example:
            # From ring_sums implementation:
            p0, p_rings = polargrid.as_rings(stationary_distribution)
            # p0 is the probability of being at the center
            # p_rings is a 2D matrix of probabilities with rows indicating r and columns indicating theta
            prob_at_r = jnp.concatenate([jnp.array([p0]), p_rings.sum(axis=1)])
            # prob_at_r is a 1D array of probabilities with length n_rvals
            
        """
        return (z[0], z[1:].reshape((self.n_rvals-1, self.n_thetavals)))

    def ring_sums(self, *, z):
        """
        Returns the sum of z values for each ring of constant radius (collapses angular dimension)
        Args:
            z: z values, a jax array of length self.len
        Returns:
            1D array of summed z values with length n_rvals

        Example:
            prob_at_r = polargrid.ring_sums(stationary_distribution)
            # prob_at_r is a 1D array of probabilities with length n_rvals
        """
        p0, p_rings = self.as_rings(z=z)
        return jnp.concatenate([jnp.array([p0]), p_rings.sum(axis=1)])

    def theta_sums(self, *, z):
        """
        Returns the sum of z values for each theta of constant angle (collapses radial dimension)
        Args:
            z: z values, a jax array of length self.len
        Returns:
            1D array of summed z values with length n_thetavals
        """
        _, p_rings = self.as_rings(z=z)
        return p_rings.sum(axis=0)

    def index(self, *, r=None, theta=None):
        if r is not None and theta is not None:
            if r>self.radius or r<0:
                raise ValueError(f"r must be between 0 and {self.radius}, got: {r}")
            if r==0:
                return 0
            theta = theta % 360
            if theta < 0:
                theta += 360
            if (theta>(360-0.501*self.thetastep)):
                theta = 0.0
            return (
                1 
                + round((r-self.rstep) / self.rstep) * self.n_thetavals 
                + round(theta / self.thetastep)
                )
        else:
            raise ValueError("Both r and theta must be specified")

    def plot(self, *, z, levels=20, cmap='viridis', label='Z value', title='Polar Contour Plot'):
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
        contour_plot = ax.contourf(self.theta, self.r, z, levels=levels, cmap=cmap)
        fig.colorbar(contour_plot, ax=ax, label=label)
        ax.set_title(title)
        plt.show()

    def parts_from_linear_discriminator(self, *, a, b, c):
        """
        Not implemented for PolarGrid
        
        """
        raise NotImplementedError()

    def partition_from_rotation(self, *, angle):
        """
        returns inverse indices for rotation angle symmetry

        Args:
            angle (int): angle in degrees

            if angle is 0, returns inverse indices for continuous rotation symmetry (each ring is a partition) (in group theory: SO(2))
            otherwise, each ring is further partitioned into 360/self.thetastep tiled segments  (in group theory: cyclic group Cn where n=360/self.thetastep)

        Returns:
            jnp.ndarray: inverse indices for use in lumping/unlumping
        """
        if angle is None:
            raise ValueError("Angle must be specified, got: None")
        # test for angle==0 first, to avoid division by zero
        if angle == 0:  # continuous rotation lumps by radius; each ring is a partition
            return jnp.round(self.r/self.rstep).astype(jnp.int32)
        if angle % self.thetastep != 0:
            raise ValueError("Angle must be a multiple of the theta step, got: {}".format(angle))
        if 360 % angle != 0:
            raise ValueError("Angle must divide 360, got: {}".format(angle))
        angle_in_thetasteps = angle // self.thetastep
        parts = (1+jnp.round(
                    ((self.r-self.rstep)/self.rstep)*angle_in_thetasteps
                    +((self.theta_deg%angle)/self.thetastep))
            ).astype(jnp.int32).at[0].set(0)
        return parts    

    def partition_from_symmetry(self, *, symmetries:list[str]):
        """
        Not implemented for PolarGrid
        """
        raise NotImplementedError
    
