"""
Grid cell network models for spatial navigation.

This module implements two grid cell models:
1. GridCell2DPosition: Position-based model with hexagonal lattice structure
2. GridCell2DVelocity: Velocity-based path integration model (Burak & Fiete 2009)
"""

import brainpy.math as bm
import jax
import numpy as np

from ._base import BasicModel


class GridCell2DPosition(BasicModel):
    """
    Position-based 2D continuous-attractor grid cell network with hexagonal lattice structure.

    This network implements a twisted torus topology that generates grid cell-like
    spatial representations with hexagonal periodicity.

    The network operates in a transformed coordinate system where grid cells form
    a hexagonal lattice, enabling realistic grid field spacing and orientation.

    Args:
        length: Number of grid cells along one dimension (total = length^2). Default: 30
        tau: Membrane time constant (ms). Default: 10.0
        k: Global inhibition strength for divisive normalization. Default: 1.0
        a: Width of connectivity kernel. Determines bump width. Default: 0.8
        A: Amplitude of external input. Default: 3.0
        J0: Peak recurrent connection strength. Default: 5.0
        mapping_ratio: Controls grid spacing (larger = smaller spacing).
            Grid spacing λ = 2π / mapping_ratio. Default: 1.5
        noise_strength: Standard deviation of activity noise. Default: 0.1
        conn_noise: Standard deviation of connectivity noise. Default: 0.0
        g: Firing rate gain factor (scales to biological range). Default: 1.0

    Attributes:
        num (int): Total number of grid cells (length^2)
        x_grid, y_grid (Array): Grid cell preferred phases in [-π, π]
        value_grid (Array): Neuron positions in phase space, shape (num, 2)
        Lambda (float): Grid spacing in real space
        coor_transform (Array): Hexagonal to rectangular coordinate transform
        coor_transform_inv (Array): Rectangular to hexagonal coordinate transform
        conn_mat (Array): Recurrent connectivity matrix
        candidate_centers (Array): Grid of candidate bump centers for decoding
        r (Variable): Firing rates (shape: num)
        u (Variable): Membrane potentials (shape: num)
        center_phase (Variable): Decoded bump center in phase space (shape: 2)
        center_position (Variable): Decoded position in real space (shape: 2)
        inp (Variable): External input for tracking (shape: num)
        gc_bump (Variable): Grid cell bump activity pattern (shape: num)

    Example:
        >>> import brainpy.math as bm
        >>> from canns.models.basic import GridCell2DPosition
        >>>
        >>> bm.set_dt(1.0)
        >>> model = GridCell2DPosition(length=16, mapping_ratio=1.5)
        >>>
        >>> # Update with a 2D position
        >>> position = bm.array([0.5, 0.3])
        >>> model.update(position)
        >>>
        >>> # Access decoded position
        >>> decoded_pos = model.center_position.value
        >>> decoded_pos.shape
        (2,)

    References:
        Burak, Y., & Fiete, I. R. (2009).
        Accurate path integration in continuous attractor network models of grid cells.
        PLoS Computational Biology, 5(2), e1000291.
    """

    def __init__(
        self,
        length: int = 30,
        tau: float = 10.0,
        k: float = 1.0,
        a: float = 0.8,
        A: float = 3.0,
        J0: float = 5.0,
        mapping_ratio: float = 1.5,
        noise_strength: float = 0.1,
        conn_noise: float = 0.0,
        g: float = 1.0,
    ):
        """Initialize the simplified grid cell network."""
        self.num = length * length
        super().__init__()

        # Store parameters
        self.length = length
        self.tau = tau
        self.k = k
        self.a = a
        self.A = A
        self.J0 = J0
        self.g = g
        self.noise_strength = noise_strength
        self.conn_noise = conn_noise
        self.mapping_ratio = mapping_ratio

        # Grid spacing in real space
        self.Lambda = 2 * bm.pi / mapping_ratio

        # Coordinate transformation matrices (hexagonal <-> rectangular)
        # coor_transform maps parallelogram (60° angle) to square
        # This partitions 2D space into parallelograms, each containing one lattice of grid cells
        self.coor_transform = bm.array([[1.0, -1.0 / bm.sqrt(3.0)], [0.0, 2.0 / bm.sqrt(3.0)]])

        # coor_transform_inv maps square to parallelogram (60° angle)
        # Equivalent to: bm.array([[1.0, 1.0/2], [0.0, bm.sqrt(3.0)/2]])
        self.coor_transform_inv = bm.linalg.inv(self.coor_transform)

        # Feature space: phase coordinates in [-π, π]
        x_bins = bm.linspace(-bm.pi, bm.pi, length + 1)
        x_grid, y_grid = bm.meshgrid(x_bins[:-1], x_bins[:-1])
        self.x_grid = x_grid.reshape(-1)
        self.y_grid = y_grid.reshape(-1)

        # Neuron positions in phase space, shape (num, 2)
        self.value_grid = bm.stack([self.x_grid, self.y_grid], axis=1)
        # Scaled positions for bump template generation
        self.value_bump = self.value_grid * 4

        # Candidate centers for position decoding (disambiguates periodic grid)
        self.candidate_centers = self.make_candidate_centers(self.Lambda)

        # Build connectivity matrix with optional noise
        base_connection = self.make_connection()
        noise_connection = bm.random.randn(self.num, self.num) * conn_noise
        self.conn_mat = base_connection + noise_connection

        # Initialize state variables
        self.r = bm.Variable(bm.zeros(self.num))  # Firing rates
        self.u = bm.Variable(bm.zeros(self.num))  # Membrane potentials
        self.inp = bm.Variable(bm.zeros(self.num))  # External input (for tracking)
        self.center_phase = bm.Variable(bm.zeros(2))  # Decoded bump center (phase)
        self.center_position = bm.Variable(bm.zeros(2))  # Decoded position (real space)
        self.gc_bump = bm.Variable(bm.zeros(self.num))  # Bump activity pattern

    def make_connection(self):
        """
        Generate recurrent connectivity matrix with 2D Gaussian kernel.

        Uses hexagonal lattice geometry via coordinate transformation.
        Connection strength decays with distance in transformed space.

        Returns:
            Array of shape (num, num): Recurrent connectivity matrix
        """

        @jax.vmap
        def kernel(v):
            # v: (2,) location in (x,y) phase space
            d = self.calculate_dist(v - self.value_grid)  # (N,) distances
            return (
                (self.J0 / self.g)
                * bm.exp(-0.5 * bm.square(d / self.a))
                / (bm.sqrt(2.0 * bm.pi) * self.a)
            )

        return kernel(self.value_grid)  # (N, N)

    def calculate_dist(self, d):
        """
        Calculate Euclidean distance after hexagonal coordinate transformation.

        Applies periodic boundary conditions and transforms displacement vectors
        from phase space to hexagonal lattice coordinates.

        Args:
            d: Displacement vectors in phase space, shape (..., 2)

        Returns:
            Array of shape (...,): Euclidean distances in hexagonal space
        """
        # Apply periodic boundary conditions
        d = self.handle_periodic_condition(d)
        # Transform to lattice axes (hex/rect)
        # This means the bump on the parallelogram lattice is a Gaussian,
        # while in the square space it is a twisted Gaussian
        dist = bm.matmul(self.coor_transform_inv, d.T).T
        return bm.sqrt(dist[:, 0] ** 2 + dist[:, 1] ** 2)

    def handle_periodic_condition(self, d):
        """
        Apply periodic boundary conditions to wrap phases into [-π, π].

        Args:
            d: Phase values (any shape with last dimension = 2)

        Returns:
            Wrapped phase values in [-π, π]
        """
        d = bm.where(d > bm.pi, d - 2.0 * bm.pi, d)
        d = bm.where(d < -bm.pi, d + 2.0 * bm.pi, d)
        return d

    def make_candidate_centers(self, Lambda):
        """
        Generate grid of candidate bump centers for decoding.

        Creates a regular lattice of potential activity bump locations
        used for disambiguating position from grid cell phases.

        Args:
            Lambda: Grid spacing in real space

        Returns:
            Array of shape (N_c*N_c, 2): Candidate centers in transformed coordinates
        """
        N_c = 32
        cc = bm.zeros((N_c, N_c, 2))

        for i in range(N_c):
            for j in range(N_c):
                cc = cc.at[i, j, 0].set((-N_c // 2 + i) * Lambda)
                cc = cc.at[i, j, 1].set((-N_c // 2 + j) * Lambda)

        cc_transformed = bm.dot(self.coor_transform_inv, cc.reshape(N_c * N_c, 2).T).T

        return cc_transformed

    def position2phase(self, position):
        """
        Convert real-space position to grid cell phase coordinates.

        Applies coordinate transformation and wraps to periodic boundaries.
        Each grid cell's preferred phase is determined by the animal's position
        on the hexagonal lattice.

        Args:
            position: Real-space coordinates, shape (2,) or (2, N)

        Returns:
            Array of shape (2,) or (2, N): Phase coordinates in [-π, π] per axis
        """
        mapped_pos = position * self.mapping_ratio
        phase = bm.matmul(self.coor_transform, mapped_pos) + bm.pi
        px = bm.mod(phase[0], 2.0 * bm.pi) - bm.pi
        py = bm.mod(phase[1], 2.0 * bm.pi) - bm.pi
        return bm.array([px, py])

    def get_unique_activity_bump(self, network_activity, animal_position):
        """
        Decode unique bump location from network activity and animal position.

        Estimates the activity bump center in phase space using population vector
        decoding, then maps it to real space and snaps to the nearest candidate
        center to resolve periodic ambiguity.

        Args:
            network_activity: Grid cell firing rates, shape (num,)
            animal_position: Current animal position for disambiguation, shape (2,)

        Returns:
            center_phase: Phase coordinates of bump center, shape (2,)
            center_position: Real-space position of bump (nearest candidate), shape (2,)
            bump: Gaussian bump template centered at center_position, shape (num,)
        """
        # Decode bump center in phase space using population vector
        exppos_x = bm.exp(1j * self.x_grid)
        exppos_y = bm.exp(1j * self.y_grid)
        activity_masked = bm.where(
            network_activity > bm.max(network_activity) * 0.1, network_activity, 0.0
        )

        center_phase = bm.zeros((2,))
        center_phase = center_phase.at[0].set(bm.angle(bm.sum(exppos_x * activity_masked)))
        center_phase = center_phase.at[1].set(bm.angle(bm.sum(exppos_y * activity_masked)))

        # Map back to real space, snap to nearest candidate center
        center_pos_residual = bm.matmul(self.coor_transform_inv, center_phase) / self.mapping_ratio
        candidate_pos_all = self.candidate_centers + center_pos_residual
        distances = bm.linalg.norm(candidate_pos_all - animal_position, axis=1)
        center_position = candidate_pos_all[bm.argmin(distances)]

        # Build Gaussian bump template
        d = bm.asarray(center_position) - self.value_bump
        dist = bm.sqrt(d[:, 0] ** 2 + d[:, 1] ** 2)
        gc_bump = self.A * bm.exp(-bm.square(dist / self.a))

        return center_phase, center_position, gc_bump

    def get_stimulus_by_pos(self, position):
        """
        Generate Gaussian stimulus centered at given position.

        Useful for previewing input patterns without calling update.

        Args:
            position: 2D position [x, y] in real space

        Returns:
            Array of shape (num,): External input pattern
        """
        position = bm.asarray(position)
        phase = self.position2phase(position)
        d = self.calculate_dist(phase - self.value_grid)
        return self.A * bm.exp(-0.5 * bm.square(d / self.a))

    def update(self, position):
        """
        Single time-step update of grid cell network dynamics.

        Updates network activity using continuous attractor dynamics with
        direct position-based external input. No adaptation or theta modulation.

        Args:
            position: Current 2D position [x, y] in real space, shape (2,)
        """
        # Convert position to array if needed
        position = bm.asarray(position)

        # 1. Decode current bump center for tracking
        center_phase, center_position, gc_bump = self.get_unique_activity_bump(
            self.r.value, position
        )
        self.center_phase.value = center_phase
        self.center_position.value = center_position
        self.gc_bump.value = gc_bump

        # 2. Calculate external input directly from position
        phase = self.position2phase(position)
        d = self.calculate_dist(phase - self.value_grid)
        Iext = self.A * bm.exp(-0.5 * bm.square(d / self.a))
        self.inp.value = Iext  # store for debugging/analysis

        # 3. Calculate recurrent input
        Irec = bm.matmul(self.conn_mat, self.r.value)

        # 4. Add activity noise
        noise = bm.random.randn(self.num) * self.noise_strength

        # 5. Update membrane potential (Euler integration)
        total_input = Irec + Iext + noise
        self.u.value += (-self.u.value + total_input) / self.tau * bm.get_dt()
        # Apply ReLU non-linearity
        self.u.value = bm.where(self.u.value > 0.0, self.u.value, 0.0)

        # 6. Calculate firing rates with divisive normalization
        u_sq = bm.square(self.u.value)
        self.r.value = self.g * u_sq / (1.0 + self.k * bm.sum(u_sq))


class GridCell2DVelocity(BasicModel):
    """
    Velocity-based grid cell network (Burak & Fiete 2009).

    This network implements path integration through velocity-modulated input
    and asymmetric connectivity. Unlike position-based models, this takes
    velocity as input and integrates it over time to track position.

    Key Features:
        - Velocity-dependent input modulation: B(v) = A * (1 + α·v·v_pref)
        - Asymmetric connectivity shifted in preferred velocity directions
        - Simple ReLU activation (not divisive normalization)
        - Healing process for proper initialization

    Args:
        length: Number of neurons along one dimension (total = length²). Default: 40
        tau: Membrane time constant. Default: 0.01
        alpha: Velocity coupling strength. Default: 0.2
        A: Baseline input amplitude. Default: 1.0
        W_a: Connection amplitude (>1 makes close surround activatory). Default: 1.5
        W_l: Spatial shift size for asymmetric connectivity. Default: 2.0
        lambda_net: Lattice constant (neurons between bump centers). Default: 15.0
        e: Controls inhibitory surround spread. Default: 1.15
            W_gamma and W_beta are computed from this and lambda_net

    Attributes:
        num (int): Total number of neurons (length²)
        positions (Array): Neuron positions in 2D lattice, shape (num, 2)
        vec_pref (Array): Preferred velocity directions (unit vectors), shape (num, 2)
        conn_mat (Array): Asymmetric connectivity matrix, shape (num, num)
        s (Variable): Neural activity/potential, shape (num,)
        r (Variable): Firing rates (ReLU of s), shape (num,)
        center_position (Variable): Decoded position in real space, shape (2,)

    Example:
        >>> import brainpy.math as bm
        >>> from canns.models.basic import GridCell2DVelocity
        >>>
        >>> bm.set_dt(5e-4)  # Small timestep for accurate integration
        >>> model = GridCell2DVelocity(length=40)
        >>>
        >>> # Healing process (recommended before simulation)
        >>> model.heal_network(num_healing_steps=50, dt_healing=1e-3)
        >>>
        >>> # Update with 2D velocity
        >>> velocity = bm.array([0.1, 0.05])  # [vx, vy]
        >>> model.update(velocity)

    References:
        Burak, Y., & Fiete, I. R. (2009).
        Accurate path integration in continuous attractor network models of grid cells.
        PLoS Computational Biology, 5(2), e1000291.
    """

    def __init__(
        self,
        length: int = 40,
        tau: float = 0.01,
        alpha: float = 0.2,
        A: float = 1.0,
        W_a: float = 1.5,
        W_l: float = 2.0,
        lambda_net: float = 15.0,
        e: float = 1.15,  # Controls inhibitory surround spread
        use_sparse: bool = False,  # Experimental: sparse matrix (for GPU)
    ):
        """Initialize the Burak & Fiete grid cell network.

        Args:
            use_sparse: Whether to use sparse matrix for connectivity (experimental).
                Default: False. Sparse matrices may be faster on GPU but slower on CPU.
                Requires brainevent library.
        """
        self.num = length * length
        super().__init__()

        # Store parameters
        self.length = length
        self.tau = tau
        self.alpha = alpha
        self.A = A
        self.W_a = W_a
        self.W_l = W_l
        self.lambda_net = lambda_net
        self.e = e
        self.use_sparse = use_sparse

        # Compute connectivity kernel parameters (from Burak & Fiete 2009)
        self.W_gamma = e * 3.0 / (lambda_net**2)  # Outer inhibitory surround
        self.W_beta = 3.0 / (lambda_net**2)  # Inner inhibitory surround

        # Create neuron positions in 2D lattice
        neuron_indices = bm.arange(0, self.num, 1)
        x_loc = neuron_indices % self.length
        y_loc = neuron_indices // self.length
        self.positions = bm.stack([x_loc, y_loc], axis=1).astype(bm.float32)

        # Generate preferred velocity directions
        self.vec_pref = self._generate_preferred_velocities()

        # Build asymmetric connectivity matrix (dense or sparse)
        self.conn_mat = self.make_connection()

        # Initialize state variables
        # Start with small random values (will be properly initialized by healing)
        self.s = bm.Variable(bm.random.rand(self.num) * 0.1)  # Neural activity/potential
        self.r = bm.Variable(bm.maximum(self.s.value, 0.0))  # Firing rates (ReLU of s)
        self.center_position = bm.Variable(bm.zeros(2))  # Decoded position

    def _generate_preferred_velocities(self):
        """
        Generate preferred velocity direction unit vectors for each neuron.

        Uses deterministic tiling with 3 basic orientations separated by 60 degrees
        to create hexagonal symmetry. Each orientation has ± directions, giving
        6 total velocity directions distributed across neurons.

        Returns:
            Array of shape (num, 2): Unit vectors indicating preferred velocity direction
        """
        # Three basic orientations for hexagonal symmetry (60° apart)
        theta1 = 0.0
        theta2 = bm.pi / 3.0
        theta3 = 2.0 * bm.pi / 3.0

        # Create deterministic pattern: assign direction based on neuron position
        # Pattern ensures hexagonal structure
        neuron_indices = bm.arange(self.num)

        # Simple deterministic assignment: cycle through 6 directions
        # Direction 0: +theta1, Direction 1: -theta1
        # Direction 2: +theta2, Direction 3: -theta2
        # Direction 4: +theta3, Direction 5: -theta3
        direction_idx = neuron_indices % 6

        # Compute theta for each neuron
        theta = bm.zeros(self.num)

        # Direction 0 and 1: theta1
        theta = bm.where(direction_idx == 0, theta1, theta)
        theta = bm.where(direction_idx == 1, theta1 + bm.pi, theta)

        # Direction 2 and 3: theta2
        theta = bm.where(direction_idx == 2, theta2, theta)
        theta = bm.where(direction_idx == 3, theta2 + bm.pi, theta)

        # Direction 4 and 5: theta3
        theta = bm.where(direction_idx == 4, theta3, theta)
        theta = bm.where(direction_idx == 5, theta3 + bm.pi, theta)

        # Convert to unit vectors
        vec_pref = bm.stack([bm.cos(theta), bm.sin(theta)], axis=1)

        return vec_pref

    def handle_periodic_condition(self, d):
        """
        Apply periodic boundary conditions to neuron position differences.

        Args:
            d: Position differences, shape (..., 2)

        Returns:
            Wrapped differences with periodic boundaries
        """
        # Wrap to [0, length)
        d = bm.where(d > self.length / 2, d - self.length, d)
        d = bm.where(d < -self.length / 2, d + self.length, d)
        return d

    def make_connection(self):
        """
        Build asymmetric connectivity matrix with spatial shifts (vectorized).

        The connectivity from neuron i to j depends on the distance between them,
        shifted by neuron i's preferred velocity direction:
            distance = |pos_j - pos_i - W_l * vec_pref_i|

        This creates asymmetric connectivity that enables velocity-driven
        bump shifts for path integration.

        Connectivity kernel:
            W_ij = W_a * (exp(-W_gamma * d²) - exp(-W_beta * d²))

        Note:
            This implementation uses JAX broadcasting for efficient computation.
            All pairwise distances are computed simultaneously, avoiding Python loops.

            If use_sparse=True, converts to brainevent.CSR sparse matrix format.
            Sparse matrices reduce memory usage for large networks but may be slower
            on CPU. They are primarily intended for GPU acceleration.

        Returns:
            Dense array of shape (num, num), or brainevent.CSR if use_sparse=True
        """
        # Vectorized computation using broadcasting (much faster than loops)
        # positions: (num, 2), vec_pref: (num, 2)

        # Compute all pairwise position differences using broadcasting
        # pos_diff[i,j] = positions[j] - positions[i] - W_l * vec_pref[i]
        # Broadcasting: (1, num, 2) - (num, 1, 2) - (num, 1, 2) → (num, num, 2)
        pos_diff = (
            self.positions[None, :, :]  # (1, num, 2) - all target positions j
            - self.positions[:, None, :]  # (num, 1, 2) - all source positions i
            - self.W_l * self.vec_pref[:, None, :]  # (num, 1, 2) - velocity-dependent shift
        )

        # Apply periodic boundary conditions
        # Wrap to [-length/2, length/2]
        pos_diff = bm.where(pos_diff > self.length / 2, pos_diff - self.length, pos_diff)
        pos_diff = bm.where(pos_diff < -self.length / 2, pos_diff + self.length, pos_diff)

        # Compute squared Euclidean distances for all pairs
        # Shape: (num, num)
        d_squared = bm.sum(pos_diff**2, axis=2)

        # Apply connectivity kernel: difference of exponentials
        # W_ij = W_a * (exp(-W_gamma * d²) - exp(-W_beta * d²))
        conn_mat = self.W_a * (bm.exp(-self.W_gamma * d_squared) - bm.exp(-self.W_beta * d_squared))

        # Convert to sparse format for large networks
        if self.use_sparse:
            try:
                import brainevent
                from scipy.sparse import csr_matrix

                # Convert to numpy and create scipy CSR
                conn_np = np.asarray(conn_mat)
                scipy_csr = csr_matrix(conn_np)

                # Create brainevent CSR (much faster for matrix-vector multiplication)
                conn_mat = brainevent.CSR(
                    (scipy_csr.data, scipy_csr.indices, scipy_csr.indptr), shape=scipy_csr.shape
                )
            except ImportError:
                print("Warning: brainevent not available, using dense matrix")
                self.use_sparse = False

        return conn_mat

    def compute_velocity_input(self, velocity):
        """
        Compute velocity-modulated input: B(v) = A * (1 + α·v·v_pref)

        Neurons whose preferred direction aligns with the velocity receive
        stronger input, creating directional modulation that drives bump shifts.

        Args:
            velocity: 2D velocity vector [vx, vy], shape (2,)

        Returns:
            Array of shape (num,): Input to each neuron
        """
        # Dot product of velocity with each neuron's preferred direction
        # Shape: (num,) = (num, 2) @ (2,)
        v_dot_vpref = bm.matmul(self.vec_pref, velocity)

        # Modulated input
        B = self.A * (1.0 + self.alpha * v_dot_vpref)

        return B

    def update(self, velocity):
        """
        Single timestep update with velocity input.

        Dynamics:
            ds/dt = (1/tau) * [-s + W·r + B(v)]
            r = ReLU(s) = max(s, 0)

        Args:
            velocity: 2D velocity [vx, vy], shape (2,)
        """
        # Convert velocity to array if needed
        velocity = bm.asarray(velocity)

        # 1. Compute velocity-modulated input
        B_v = self.compute_velocity_input(velocity)

        # 2. Compute recurrent input
        # Use @ operator for both dense and sparse matrices
        Irec = self.conn_mat @ self.r.value

        # 3. Update state variable s (Euler integration)
        ds = (-self.s.value + Irec + B_v) / self.tau * bm.get_dt()
        self.s.value += ds

        # 4. Compute firing rate with ReLU
        self.r.value = bm.maximum(self.s.value, 0.0)

    def heal_network(self, num_healing_steps=2500, dt_healing=1e-4):
        """
        Healing process to form stable activity bump before simulation (optimized).

        This process is critical for proper initialization. It relaxes the network
        to a stable attractor state through a sequence of movements:
        1. Relax with zero velocity (T=0.25s)
        2. Move in 4 cardinal directions (0°, 90°, 180°, 270°)
        3. Relax again with zero velocity (T=0.25s)

        Args:
            num_healing_steps: Total number of healing steps. Default: 2500
            dt_healing: Small timestep for healing integration. Default: 1e-4

        Note:
            This temporarily changes the global timestep. The original timestep
            is restored after healing. Uses bm.for_loop for efficient execution.
        """
        # Save current dt
        original_dt = bm.get_dt()
        bm.set_dt(dt_healing)

        # Prepare velocity sequence for all healing steps
        steps_relax = int(0.25 / dt_healing)
        steps_per_dir = int(0.125 / dt_healing)

        # Phase 1: Zero velocity (relax)
        velocities_phase1 = bm.zeros((steps_relax, 2))

        # Phase 2: Move in 4 cardinal directions
        v_norm = 0.8
        angles = bm.array([0.0, bm.pi / 5, bm.pi / 2, -bm.pi / 5])
        velocities_phase2 = []
        for angle in angles:
            direction = v_norm * bm.stack([bm.cos(angle), bm.sin(angle)])
            velocities_phase2.append(bm.tile(direction, (steps_per_dir, 1)))
        velocities_phase2 = bm.concatenate(velocities_phase2, axis=0)

        # Phase 3: Zero velocity (relax again)
        velocities_phase3 = bm.zeros((steps_relax, 2))

        # Concatenate all phases
        velocities = bm.concatenate(
            [velocities_phase1, velocities_phase2, velocities_phase3], axis=0
        )

        # Define step function for bm.for_loop
        def healing_step(i, vel):
            self.update(vel)

        # Run all healing steps with bm.for_loop
        bm.for_loop(healing_step, (bm.arange(len(velocities)), velocities), progress_bar=True)

        # Restore original dt
        bm.set_dt(original_dt)

    def decode_position_lsq(self, activity_history, velocity_history):
        """
        Decode position using velocity integration (simple method).

        For proper position decoding from neural activity, a more sophisticated
        method would fit the activity to spatial basis functions. For now,
        we use velocity integration as ground truth and compute error metrics.

        Args:
            activity_history: Neural activity over time, shape (T, num)
            velocity_history: Velocity over time, shape (T, 2)

        Returns:
            decoded_positions: Integrated positions, shape (T, 2)
            r_squared: R² score (comparing integrated vs true positions if available)
        """
        velocity_np = np.asarray(velocity_history)

        # Integrate velocity to get position
        dt = bm.get_dt()
        # Start from origin [0, 0] and integrate
        integrated_pos = np.zeros_like(velocity_np)
        integrated_pos[0] = velocity_np[0] * dt  # First step
        for i in range(1, len(velocity_np)):
            integrated_pos[i] = integrated_pos[i - 1] + velocity_np[i] * dt

        # For R² computation, we would need ground truth position
        # For now, assume perfect integration gives R²=1.0
        # In practice, you'd compare to external position measurements
        r_squared = 1.0  # Placeholder - velocity integration is our reference

        return integrated_pos, r_squared

    def decode_position_from_activity(self, activity):
        """
        Decode position from neural activity using population vector method.

        This method analyzes the activity bump to determine the network's
        internal representation of position. Currently simplified.

        Args:
            activity: Neural activity, shape (num,)

        Returns:
            position: Decoded 2D position, shape (2,)
        """
        # Find peak activity location
        peak_idx = np.argmax(activity)

        # Convert neuron index to grid coordinates
        x_idx = peak_idx % self.length
        y_idx = peak_idx // self.length

        # Map to position space (simplified - assumes lattice spacing)
        # This is a placeholder; proper implementation would use phase decoding
        position = np.array([x_idx / self.length, y_idx / self.length]) * 2.2

        return position

    @staticmethod
    def track_blob_centers(activities, length):
        """
        Track blob centers using Gaussian filtering and thresholding.

        This is the robust method from Burak & Fiete 2009 reference implementation
        that achieves R² > 0.99 for path integration quality.

        Args:
            activities: Neural activities, shape (T, num)
            length: Grid size (e.g., 40 for 40×40 grid)

        Returns:
            centers: Blob centers in neuron coordinates, shape (T, 2)

        Example:
            >>> activities = np.array([...])  # (T, 1600) for 40×40 grid
            >>> centers = GridCell2DVelocity.track_blob_centers(activities, length=40)
            >>> # centers.shape == (T, 2)
        """
        from scipy.ndimage import center_of_mass, gaussian_filter, label

        T = len(activities)
        n = length

        # Reshape and apply Gaussian smoothing (per-frame to avoid axes parameter)
        activities_2d = activities.reshape(T, n, n)
        smoothed = np.array([gaussian_filter(activities_2d[t], sigma=1) for t in range(T)])

        # Adaptive thresholding
        thresholds = smoothed.mean(axis=(1, 2)) + 0.5 * smoothed.std(axis=(1, 2))
        binary_images = smoothed > thresholds[:, None, None]

        # Track centers
        centers = []
        for i in range(T):
            labeled, num_features = label(binary_images[i])

            if num_features > 0:
                blob_centers = np.array(
                    center_of_mass(binary_images[i], labeled, range(1, num_features + 1))
                )

                if blob_centers.ndim == 1:
                    blob_centers = blob_centers.reshape(1, -1)

                blob_centers = blob_centers[:, [1, 0]]  # Swap x,y
                dist = np.linalg.norm(blob_centers - n / 2, axis=1)
                best_center = blob_centers[np.argmin(dist)]
            else:
                best_center = centers[-1] if centers else np.array([n / 2, n / 2])

            centers.append(best_center)

        return np.array(centers)
