
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import scipy.sparse
from dataclasses import dataclass, field
from typing import Optional


# Import from core
from .core import constants


@dataclass
class PartitionMetadata:
    """Metadata about partition structure for optimization.
    
    Attributes:
        lazy_lumpable: If True, suggests checking lazy lumping (not guaranteed to work)
        primary_group_size: Size of most groups (None if variable)
        exceptions: Mapping of group_id to size for non-primary groups
        source: How partition was created (e.g., 'rotation_120deg')
    """
    lazy_lumpable: bool = False
    primary_group_size: Optional[int] = None
    exceptions: dict[int, int] = field(default_factory=dict)
    source: str = ""


class Partition:
    """Partition representation with metadata for optimization.
    
    A partition divides n states into k groups. This class wraps the partition
    array with metadata that enables optimization of lumping operations.
    
    Attributes:
        inverse_indices: Array mapping each state to its group (0 to k-1)
        metadata: Optional metadata about partition structure
        k: Number of groups in partition
    
    Note:
        To use partition in array operations, access `.inverse_indices` directly:
        ```python
        partition = grid.partition_from_rotation(angle=120)
        unique_groups = jnp.unique(partition.inverse_indices)
        ```
    """
    
    def __init__(self, inverse_indices: jnp.ndarray, metadata: Optional[PartitionMetadata] = None):
        """Initialize a Partition.
        
        Args:
            inverse_indices: Array of shape (n,) mapping states to groups
            metadata: Optional metadata for optimization
        """
        self.inverse_indices = jnp.asarray(inverse_indices, dtype=jnp.int32)
        self.metadata = metadata or PartitionMetadata()
        self.k = int(self.inverse_indices.max()) + 1
    
    def __jax_tree_flatten__(self):
        """Make Partition a JAX pytree for JIT compatibility."""
        children = (self.inverse_indices,)
        aux_data = (self.metadata, self.k)
        return children, aux_data
    
    @classmethod
    def __jax_tree_unflatten__(cls, aux_data, children):
        """Reconstruct Partition from pytree."""
        metadata, k = aux_data
        inverse_indices, = children
        return cls(inverse_indices, metadata)
    
    def __len__(self):
        """Return number of states."""
        return len(self.inverse_indices)
    
    def max(self):
        """Return maximum value in inverse_indices (for backward compatibility)."""
        return self.inverse_indices.max()


# Register Partition as a JAX pytree
jax.tree_util.register_pytree_node(
    Partition,
    Partition.__jax_tree_flatten__,
    Partition.__jax_tree_unflatten__
)



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 = constants.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*constants.EPSILON*jnp.abs(min_z)
        max_z = float(z.max())
        max_z_mask = jnp.abs(z-max_z) <= 2*constants.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=constants.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), symmetry_type: str = "") -> 'Partition':
        """
        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))
            symmetry_type: Type of symmetry (e.g., "reflect_x", "reflect_y", "swap_xy")
        
        Returns:
            Partition: Partition object grouping points by discriminant value
        
        Notes:
            - Used internally for efficient symmetry partitioning
            - Optimized for reflection symmetries on regular grids
            - For reflections, most groups have size 2 (reflected pairs)
            - Points on reflection axis form singleton groups (size 1)
        """
        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]
        
        # Count group sizes to identify singleton exceptions (points on reflection axis)
        group_sizes = jnp.bincount(inverse_indices, length=int(inverse_indices.max()) + 1)
        exceptions = {int(g): int(size) for g, size in enumerate(group_sizes) if size == 1}
        
        metadata = PartitionMetadata(
            lazy_lumpable=True,
            primary_group_size=2,  # Most groups are reflected pairs
            exceptions=exceptions,  # Singleton groups on reflection axis
            source=symmetry_type if symmetry_type else "linear_discriminator"
        )
        return Partition(inverse_indices, metadata)

    def partition_from_symmetry(
        self,
        symmetries: list,
        tolerance: float = 1e-6
    ) -> 'Partition':
        """
        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:
            Partition: Partition object grouping symmetric grid points
        
        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
        
        # Empty symmetries: identity partition
        if len(symmetries) == 0:
            inverse_indices = jnp.arange(n_states, dtype=jnp.int32)
            metadata = PartitionMetadata(
                lazy_lumpable=False,
                primary_group_size=1,
                exceptions={},
                source="identity"
            )
            return Partition(inverse_indices, metadata)
        
        # 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),
                        symmetry_type=sym
                    )
                
                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),
                        symmetry_type=sym
                    )
                
                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|
                        symmetry_type=sym
                    )
        
        # 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
        inverse_indices = jnp.array(labels, dtype=jnp.int32)
        
        # Create descriptive source string by concatenating symmetry names
        symmetry_names = []
        for sym in symmetries:
            if isinstance(sym, str):
                symmetry_names.append(sym)
            elif isinstance(sym, tuple) and sym[0] == 'rotate':
                # Format: rotate_cx_cy_degdeg
                symmetry_names.append(f"rotate_{sym[1]}_{sym[2]}_{sym[3]}deg")
        source = "_".join(symmetry_names)
        
        metadata = PartitionMetadata(
            lazy_lumpable=True,  # Hint (not guaranteed)
            primary_group_size=None,  # Variable - depends on symmetry combination
            exceptions={},
            source=source
        )
        return Partition(inverse_indices, metadata)


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
            inverse_indices = jnp.round(self.r/self.rstep).astype(jnp.int32)
            metadata = PartitionMetadata(
                lazy_lumpable=True,
                primary_group_size=None,  # Variable: each ring has different number of states
                exceptions={},  # All groups are different sizes
                source='rotation_continuous'
            )
            return Partition(inverse_indices, metadata)
        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
        inverse_indices = (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)
        
        # Create metadata for discrete rotation symmetry
        primary_group_size = 360 // angle  # Number of states per group (except origin)
        metadata = PartitionMetadata(
            lazy_lumpable=True,  # Suggests lazy lumping (not guaranteed)
            primary_group_size=primary_group_size,
            exceptions={0: 1},  # Origin group has 1 state
            source=f'rotation_{angle}deg'
        )
        return Partition(inverse_indices, metadata)    

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