"""Systematic rate map sampling for grid cell analysis.

This module provides functions for computing neural rate maps through systematic
spatial sampling instead of trajectory-based methods. This approach ensures:
- 100% spatial coverage (no gaps)
- Uniform sampling density across all locations
- Stable continuous attractor dynamics preservation
- Higher quality grid scores for grid cell analysis

The systematic sampling method samples the entire 2D spatial grid in a structured way:
1. Horizontal sweep: Move from left to right, saving network state at each position
2. Vertical sampling: From each horizontal position state, scan all vertical positions
3. Assembly: Combine responses into complete (resolution × resolution) rate map

This matches the approach used in the reference Burak & Fiete (2009) implementation
and produces much cleaner grid cell firing fields compared to trajectory-based methods.

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

Example:
    >>> from canns.models.basic import GridCell2DVelocity
    >>> from canns.analyzer.metrics.systematic_ratemap import compute_systematic_ratemap
    >>>
    >>> # Initialize and heal network
    >>> model = GridCell2DVelocity(length=40)
    >>> model.heal_network(num_healing_steps=5000)
    >>>
    >>> # Compute systematic rate map
    >>> ratemap = compute_systematic_ratemap(
    ...     model,
    ...     box_width=2.2,
    ...     box_height=2.2,
    ...     resolution=30,
    ...     speed=0.3,
    ... )
    >>> ratemap.shape  # (resolution, resolution, num_neurons)
    (30, 30, 1600)
"""

import brainpy.math as bm
import numpy as np
from numba import njit
from tqdm import tqdm

__all__ = [
    "compute_systematic_ratemap",
]


# ============================================================================
# Numba-Optimized Helper Functions
# ============================================================================


@njit
def _compute_velocities(positions: np.ndarray, dt: float) -> np.ndarray:
    """Compute velocities from position trajectory.

    Args:
        positions: Position array of shape (T, 2) with (x, y) coordinates
        dt: Time step size

    Returns:
        Velocity array of shape (T-1, 2) with (vx, vy) velocities
    """
    # Manual diff implementation for Numba compatibility
    T = positions.shape[0]
    velocities = np.zeros((T - 1, 2))
    for i in range(T - 1):
        velocities[i] = positions[i + 1] - positions[i]
    return velocities / dt


@njit
def _create_vertical_velocities(num_steps: int, num_positions: int, speed: float) -> np.ndarray:
    """Pre-allocate vertical velocity array.

    Args:
        num_steps: Number of vertical steps
        num_positions: Number of horizontal positions
        speed: Movement speed in m/s

    Returns:
        Velocity array of shape (num_steps, num_positions, 2) with upward velocities
    """
    batch_vel = np.zeros((num_steps, num_positions, 2))
    batch_vel[:, :, 1] = speed  # Upward velocity
    return batch_vel


@njit
def _downsample_indices(total_length: int, target_size: int) -> np.ndarray:
    """Compute downsampling indices.

    Args:
        total_length: Total length of array to downsample
        target_size: Desired output size

    Returns:
        Array of indices to extract from original array
    """
    ratio = total_length // target_size
    indices = np.arange(target_size) * ratio
    # Ensure we don't go out of bounds
    indices = np.minimum(indices, total_length - 1)
    return indices


@njit
def _downsample_activities(
    activities: np.ndarray, downsample_ratio: int, resolution: int
) -> np.ndarray:
    """Downsample vertical sweep activities to spatial bins.

    Args:
        activities: Activity array of shape (num_steps, batch_size, num_neurons)
        downsample_ratio: Ratio for downsampling timesteps to spatial bins
        resolution: Target spatial resolution

    Returns:
        Downsampled array of shape (batch_size, resolution, num_neurons)
    """
    num_steps, batch_size, num_neurons = activities.shape
    result = np.zeros((batch_size, resolution, num_neurons))

    for i in range(batch_size):
        for y_idx in range(resolution):
            t_idx = y_idx * downsample_ratio
            if t_idx < num_steps:
                result[i, y_idx, :] = activities[t_idx, i, :]

    return result


# ============================================================================
# Public API
# ============================================================================


def compute_systematic_ratemap(
    model,
    box_width: float = 2.2,
    box_height: float = 2.2,
    resolution: int = 100,
    speed: float = 0.5,
    num_batches: int = 10,
    verbose: bool = True,
) -> np.ndarray:
    """Compute rate maps by systematically sampling the entire spatial grid.

    This function implements the systematic spatial sampling approach used in
    the reference Burak & Fiete (2009) grid cell model. Instead of relying on
    trajectory coverage, it systematically samples every point in the spatial
    grid to ensure complete and uniform coverage.

    The algorithm works in three steps:
    1. **Horizontal sweep**: Move from left to right along the bottom edge,
       saving the network state at each horizontal position
    2. **Vertical sampling**: From each saved horizontal state, move upward
       sampling neural activity at each vertical position
    3. **Assembly**: Combine all sampled activities into a complete rate map

    IMPORTANT - Sampling Method:
    This is NOT a continuous trajectory. Instead, it uses state restoration:
    - Horizontal sweep (Step 1): Continuous movement via velocity input ✓
    - Vertical sampling (Step 2): For each x-position, the network state is
      RESTORED (jumped) to the saved state, then moved continuously upward

    Why use state restoration instead of continuous scanning?
    - Ensures uniform sampling density across all spatial locations
    - Computationally efficient (avoids redundant back-and-forth movement)
    - Maintains continuous attractor dynamics within each vertical sweep

    This is a standard technique for spatial analysis of grid cells, trading
    strict continuity for uniform spatial coverage. The bump activity remains
    stable because each restored state comes from continuous movement.

    CRITICAL: This function assumes the model has already been healed via
    `model.heal_network()` to establish a stable attractor state. The healed
    state is automatically preserved - DO NOT call `model.reset()` before or
    after this function as it will destroy the healed state.

    Args:
        model: Grid cell model (must have been healed via model.heal_network())
        box_width: Arena width in meters (default: 2.2)
        box_height: Arena height in meters (default: 2.2)
        resolution: Number of spatial bins (creates resolution×resolution grid)
            (default: 100)
        speed: Movement speed in m/s (default: 0.5)
        num_batches: Number of batches to split computation for memory efficiency
            (default: 10)
        verbose: Print progress information (default: True)

    Returns:
        ratemap: Rate map array of shape (resolution, resolution, num_neurons)
            containing the neural activity at each spatial location

    Example:
        >>> import brainpy.math as bm
        >>> from canns.models.basic import GridCell2DVelocity
        >>> from canns.analyzer.metrics.systematic_ratemap import compute_systematic_ratemap
        >>>
        >>> # Dummy model (small for quick runs)
        >>> bm.set_dt(5e-4)
        >>> model = GridCell2DVelocity(length=10, alpha=0.1, lambda_net=17.0)
        >>> model.heal_network(num_healing_steps=50, dt_healing=1e-4)
        >>>
        >>> ratemap = compute_systematic_ratemap(
        ...     model,
        ...     box_width=1.0,
        ...     box_height=1.0,
        ...     resolution=5,
        ...     speed=0.3,
        ...     num_batches=1,
        ...     verbose=False,
        ... )
        >>> print(ratemap.shape)
        (5, 5, 100)

    Notes:
        - Expected grid scores: 0.4-0.8 (much higher than trajectory-based methods)
        - Computation time scales with resolution^2 and num_neurons
        - For faster computation with slight quality tradeoff, reduce resolution
        - The healed state is automatically preserved; no need to save/restore it

    References:
        Burak & Fiete (2009). PLoS Computational Biology, 5(2), e1000291.
        See `.ref/grid_cells_burak_fiete/evaluate_grid.ipynb` for reference implementation.
    """
    # Get simulation parameters
    dt = model.dt if hasattr(model, "dt") else bm.get_dt()
    dx = speed * dt

    if verbose:
        print("Computing systematic rate maps...")
        print(f"  Resolution: {resolution}×{resolution}")
        print(f"  Arena: {box_width}×{box_height}m")
        print(f"  Speed: {speed} m/s, dt: {dt}s, dx: {dx}m")

    # ========================================================================
    # Step 1: Horizontal sweep to establish initial states
    # ========================================================================

    # Create horizontal trajectory from left to right at bottom edge
    c = np.arange(-box_width / 2, box_width / 2, dx)
    y_bottom = -box_height / 2
    horizontal_pos = np.stack([c, np.ones_like(c) * y_bottom], axis=1)

    # Compute velocities for horizontal movement (Numba-optimized)
    horizontal_vel = _compute_velocities(horizontal_pos, dt)

    if verbose:
        print(f"  Step 1: Horizontal sweep ({len(horizontal_vel)} steps)")

    # Run horizontal sweep to move network state
    def run_step(i, vel):
        model(vel)
        return model.r.value

    _ = bm.for_loop(
        run_step,
        (bm.arange(len(horizontal_vel)), bm.asarray(horizontal_vel)),
        progress_bar=False,
    )

    # ========================================================================
    # Step 2: Record states at horizontal positions
    # ========================================================================

    # Save the healed state to restore before horizontal sweep
    # CRITICAL: Do NOT use model.reset() as it destroys the healed state!
    healed_state = model.s.value.copy()
    healed_r = model.r.value.copy()

    if verbose:
        print(f"  Step 2: Recording states at {resolution} horizontal positions...")

    # Restore healed state before recording horizontal states
    model.s.value = healed_state
    model.r.value = healed_r

    # Downsample to resolution positions
    ratio = int(len(c) / resolution)
    horizontal_states_s = []
    horizontal_states_r = []

    # Re-run horizontal sweep, saving states at downsampled positions
    for i in range(len(horizontal_vel)):
        model(horizontal_vel[i])
        if i % ratio == 0 and len(horizontal_states_s) < resolution:
            horizontal_states_s.append(model.s.value.copy())
            horizontal_states_r.append(model.r.value.copy())

    # Ensure we have exactly resolution states
    horizontal_states_s = horizontal_states_s[:resolution]
    horizontal_states_r = horizontal_states_r[:resolution]
    horizontal_states_s = bm.asarray(horizontal_states_s)
    horizontal_states_r = bm.asarray(horizontal_states_r)

    # ========================================================================
    # Step 3: Vertical sampling at each horizontal position
    # ========================================================================

    if verbose:
        print(f"  Step 3: Vertical sampling in {num_batches} batches...")

    # Create vertical movement velocities (Numba-optimized)
    vertical_distance = box_height
    num_vertical_steps = int(vertical_distance / dx)
    batch_vel = _create_vertical_velocities(num_vertical_steps, resolution, speed)

    # Initialize rate map
    ratemap = np.zeros((resolution, resolution, model.num))

    # Process in batches for memory efficiency
    samples_per_batch = resolution // num_batches

    for batch_idx in tqdm(range(num_batches), desc="Batching", disable=not verbose):
        start_idx = batch_idx * samples_per_batch
        end_idx = min((batch_idx + 1) * samples_per_batch, resolution)
        batch_size = end_idx - start_idx

        # Get initial states for this batch
        batch_initial_states_s = horizontal_states_s[start_idx:end_idx]
        batch_initial_states_r = horizontal_states_r[start_idx:end_idx]

        # Run vertical sweeps for this batch
        batch_activities = np.zeros((num_vertical_steps, batch_size, model.num))

        for i in range(batch_size):
            # Set initial state from horizontal sweep (restore both s and r)
            model.s.value = batch_initial_states_s[i]
            model.r.value = batch_initial_states_r[i]

            # Run vertical sweep with bm.for_loop for better performance
            def run_vertical_step(t, vel):
                model(vel)
                return model.r.value

            velocities_i = bm.asarray(batch_vel[:, i])
            activities = bm.for_loop(
                run_vertical_step,
                (bm.arange(num_vertical_steps), velocities_i),
                progress_bar=False,
            )
            # Keep as BrainPy array until batching is done (avoid premature conversion)
            batch_activities[:, i, :] = activities

        # Downsample to resolution using Numba-optimized function
        downsample_ratio = num_vertical_steps // resolution
        downsampled = _downsample_activities(batch_activities, downsample_ratio, resolution)

        # Store in ratemap
        ratemap[start_idx:end_idx, :, :] = downsampled

    if verbose:
        print(f"  ✓ Rate maps computed: shape={ratemap.shape}")
        print(f"    Activity range: [{ratemap.min():.4f}, {ratemap.max():.4f}]")

    return ratemap
