src/core/network.py

"""Sparse 5D spiking neural network for MHRN.

This module defines the NeuralNetwork class, which manages:
- Neurons in a 5D spatial grid
- Synaptic connections with STDP plasticity
- Spike propagation with configurable delays
- Event queues for temporal processing
- Input/output layer management
- Post-step hooks for observers

The network follows a tick-based simulation model where:
- Each step() processes exactly one tick (1ms)
- External currents are applied before synaptic currents
- Spike events are queued with delays
- Post-step hooks run after each tick
"""

from __future__ import annotations

import logging
import random
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, cast

from .neuron import Neuron, NeuronConfig, NeuronType, create_neuron
from .spatial_index import (
    DIM_NAMES,
    Coord5D,
    Dim5D,
    iter_neighbour_coords,
    pack_coords,
    unpack_coords,
    validate_coord_in_dims,
    validate_dims,
)
from .synapse import Synapse, SynapseConfig, create_synapse

_LOGGER = logging.getLogger(__name__)


# ============================================================================
# Type Aliases
# ============================================================================

PostStepHook = Callable[["StepResult"], None]
"""Callback function type for post-step hooks."""

ConfigDict = dict[str, object]
"""Type alias for a configuration dictionary passed to NeuralNetwork.__init__."""


# ============================================================================
# Configuration Classes
# ============================================================================


@dataclass(frozen=True, slots=True)
class SimulationConfig:
    """Configuration for simulation parameters."""

    dt_ms: float = 1.0
    max_delay: int = 5
    debug_invariants: bool = False

    def __post_init__(self) -> None:
        if self.dt_ms != 1.0:
            raise ValueError("dt_ms must be 1.0 for the reference core")
        if self.max_delay < 1:
            raise ValueError("max_delay must be >= 1")


@dataclass(frozen=True, slots=True)
class TopologyConfig:
    """Configuration for network topology."""

    allow_self_connections: bool = False
    allow_parallel_connections: bool = False


@dataclass(frozen=True, slots=True)
class NetworkConfig:
    """Configuration for network parameters."""

    weight_min: float = 0.0
    weight_max: float = 0.5
    initial_connections_per_neuron: int = 10
    neighbour_radius: float = 5.0


@dataclass(frozen=True, slots=True)
class Brain5DConfig:
    """Complete configuration for MHRN network."""

    dimensions: Dim5D
    simulation: SimulationConfig = field(default_factory=SimulationConfig)
    topology: TopologyConfig = field(default_factory=TopologyConfig)
    network: NetworkConfig = field(default_factory=NetworkConfig)
    neuron: NeuronConfig = field(default_factory=NeuronConfig)
    synapse: SynapseConfig = field(default_factory=SynapseConfig)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Brain5DConfig:
        """Create config from dictionary (backward compatibility)."""
        dims = data.get("dimensions")
        if not dims or len(dims) != 5:
            raise ValueError("dimensions must be a tuple/list of 5 ints")

        sim = data.get("simulation", {})
        topo = data.get("topology", {})
        net = data.get("network", {})
        neuron_data = cast(dict[str, Any], data.get("neuron", {}))
        energy_data = cast(dict[str, Any], data.get("energy", {}))
        neuron_payload = dict(neuron_data)
        neuron_payload.setdefault("dt_ms", float(sim.get("dt_ms", 1.0)))
        neuron_payload["spike_cost"] = float(
            energy_data.get("spike_cost", neuron_payload.get("spike_cost", 0.001))
        )
        neuron_payload["resting_energy"] = float(
            energy_data.get("initial", neuron_payload.get("resting_energy", 1.0))
        )

        return cls(
            dimensions=tuple(dims),
            simulation=SimulationConfig(
                dt_ms=float(sim.get("dt_ms", 1.0)),
                max_delay=int(sim.get("max_delay", 5)),
                debug_invariants=bool(sim.get("debug_invariants", False)),
            ),
            topology=TopologyConfig(
                allow_self_connections=bool(topo.get("allow_self_connections", False)),
                allow_parallel_connections=bool(
                    topo.get("allow_parallel_connections", False)
                ),
            ),
            network=NetworkConfig(
                weight_min=float(net.get("weight_min", 0.0)),
                weight_max=float(net.get("weight_max", 0.5)),
                initial_connections_per_neuron=int(
                    net.get("initial_connections_per_neuron", 10)
                ),
                neighbour_radius=float(net.get("neighbour_radius", 5.0)),
            ),
            neuron=NeuronConfig.from_dict(neuron_payload),
            synapse=SynapseConfig(
                a_plus=float(data.get("stdp", {}).get("a_plus", 0.1)),
                a_minus=float(data.get("stdp", {}).get("a_minus", 0.12)),
                tau_plus=float(data.get("stdp", {}).get("tau_plus", 20.0)),
                tau_minus=float(data.get("stdp", {}).get("tau_minus", 20.0)),
                w_min=float(net.get("weight_min", 0.0)),
                w_max=float(net.get("weight_max", 0.5)),
                enable_triplet=bool(data.get("stdp", {}).get("enable_triplet", False)),
                enable_metaplasticity=bool(
                    data.get("stdp", {}).get("enable_metaplasticity", False)
                ),
            ),
        )


# ============================================================================
# Event Classes
# ============================================================================


@dataclass(slots=True)
class SpikeEvent:
    """A queued spike event for future delivery.

    Attributes:
        source_id: ID of the neuron that fired the spike.
        target_id: ID of the target neuron.
        weight: Synaptic weight for this event.
        delivery_tick: Tick at which this event should be delivered.
    """

    source_id: int
    target_id: int
    weight: float
    delivery_tick: int


@dataclass(slots=True)
class StepResult:
    """Result of a single network step.

    Attributes:
        tick: The tick number that was processed.
        spike_ids: IDs of neurons that spiked this tick.
        output_spike_ids: IDs of output neurons that spiked.
        spikes_this_tick: Number of spikes in this tick.
        total_spikes: Total spikes since network creation.
        delivered_events: Number of events delivered this tick.
        queued_events: Number of events currently queued.
        external_injection_count: Number of neurons with external current.
        external_total_current: Sum of all external currents.
        synaptic_current_targets: Number of neurons receiving synaptic current.
        mean_v: Mean membrane potential across all neurons.
        min_v: Minimum membrane potential.
        max_v: Maximum membrane potential.
        mean_energy: Mean energy across all neurons.
        core_step_ms: Time taken for the step in milliseconds.
        neuron_activity: Dictionary mapping neuron_id to spike flag.
        total_synapses: Total synapses in the network.
    """

    tick: int = 0
    spike_ids: tuple[int, ...] = ()
    output_spike_ids: tuple[int, ...] = ()
    spikes_this_tick: int = 0
    total_spikes: int = 0
    delivered_events: int = 0
    queued_events: int = 0
    external_injection_count: int = 0
    external_total_current: float = 0.0
    synaptic_current_targets: int = 0
    mean_v: float = 0.0
    min_v: float = 0.0
    max_v: float = 0.0
    mean_energy: float = 0.0
    core_step_ms: float = 0.0
    neuron_activity: dict[int, bool] = field(default_factory=dict[int, bool])
    total_synapses: int = 0
    dirty_neuron_ids: tuple[int, ...] = ()
    dirty_synapse_ids: tuple[tuple[int, int], ...] = ()

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "tick": self.tick,
            "spike_ids": list(self.spike_ids),
            "output_spike_ids": list(self.output_spike_ids),
            "spikes_this_tick": self.spikes_this_tick,
            "total_spikes": self.total_spikes,
            "delivered_events": self.delivered_events,
            "queued_events": self.queued_events,
            "external_injection_count": self.external_injection_count,
            "external_total_current": self.external_total_current,
            "synaptic_current_targets": self.synaptic_current_targets,
            "mean_v": self.mean_v,
            "min_v": self.min_v,
            "max_v": self.max_v,
            "mean_energy": self.mean_energy,
            "core_step_ms": self.core_step_ms,
            "total_synapses": self.total_synapses,
            "dirty_neuron_ids": list(self.dirty_neuron_ids),
            "dirty_synapse_ids": [list(value) for value in self.dirty_synapse_ids],
        }


# ============================================================================
# NeuralNetwork Class
# ============================================================================


class NeuralNetwork:
    """Sparse 5D spiking neural network.

    The network manages neurons, synapses, and spike propagation with
    configurable delays. It follows a tick-based simulation model:

    Tick semantics:
    1. External currents queued for the current tick are applied.
    2. Spike events with delivery_tick == current_tick are delivered.
    3. Neurons are updated with combined external + synaptic currents.
    4. Spikes generated in this tick are queued for future ticks.
    5. Post-step hooks run after the core step is complete.
    6. current_tick increments by 1 after the step.

    The network uses a 5D spatial grid for neuron placement, with
    distance-based connectivity and boundary layers for input/output.

    Example:
        >>> config = Brain5DConfig(dimensions=(10,10,10,10,10))
        >>> network = NeuralNetwork(config, rng)
        >>> network.add_neuron((1,2,3,4,5))
        >>> network.initialize_random_connections(10, 2.0)
        >>> result = network.step()
    """

    def __init__(
        self,
        config: Brain5DConfig | dict[str, Any] | None = None,
        rng: random.Random | None = None,
    ) -> None:
        """Initialize the neural network.

        Args:
            config: Network configuration. Can be a Brain5DConfig instance
                or a dict for backward compatibility.
            rng: Random number generator. Created with default seed if not provided.

        Raises:
            ValueError: If dimensions are invalid or configuration is malformed.
        """
        # Parse configuration
        if config is None:
            config = Brain5DConfig(dimensions=(50, 50, 50, 50, 50))
        elif isinstance(config, dict):
            config = Brain5DConfig.from_dict(config)

        self.config = config
        self.dimensions: Dim5D = config.dimensions
        self.sim_config = config.simulation
        self.topology_config = config.topology
        self.network_config = config.network
        self.neuron_config = config.neuron
        self.synapse_config = config.synapse

        # Validate dimensions
        validate_dims(self.dimensions)

        # Random number generator
        self.rng = rng or random.Random()

        # Core data structures
        self.neurons: dict[int, Neuron] = {}
        self.synapses: dict[int, list[Synapse]] = {}
        self.in_degree: dict[int, int] = {}

        # Event queue (circular buffer)
        self.max_delay = self.sim_config.max_delay
        self.event_slots: list[list[SpikeEvent]] = [
            [] for _ in range(self.max_delay + 1)
        ]
        self._queued_event_count = 0
        self._synapse_count = 0

        # State
        self.current_tick = 0
        self.total_spikes = 0
        self.total_events_processed = 0
        self.pending_currents: dict[int, float] = {}

        # Input/output cells
        self.input_cells: set[int] = set()
        self.output_cells: set[int] = set()

        # Post-step hooks
        self._post_step_hooks: list[PostStepHook] = []

        # Performance tracking
        self._step_count = 0
        self._dirty_neuron_ids: set[int] = set()
        self._dirty_synapse_ids: set[tuple[int, int]] = set()

    # ========================================================================
    # Configuration Access
    # ========================================================================

    @property
    def debug_invariants(self) -> bool:
        """Check if debug invariants are enabled."""
        return self.sim_config.debug_invariants

    @property
    def allow_self_connections(self) -> bool:
        """Check if self-connections are allowed."""
        return self.topology_config.allow_self_connections

    @property
    def allow_parallel_connections(self) -> bool:
        """Check if parallel connections are allowed."""
        return self.topology_config.allow_parallel_connections

    # ========================================================================
    # Neuron Management
    # ========================================================================

    def add_neuron(
        self,
        coord: Coord5D,
        neuron_type: NeuronType = NeuronType.REGULAR_SPIKING,
        **kwargs: Any,
    ) -> int:
        """Add a neuron at the specified 5D coordinate.

        Args:
            coord: 5D coordinate for the neuron.
            neuron_type: Type of neuron (affects default parameters).
            **kwargs: Additional parameters passed to create_neuron.

        Returns:
            The unique neuron ID (packed coordinate).

        Raises:
            ValueError: If the coordinate is outside dimensions.
            KeyError: If a neuron already exists at this coordinate.
        """
        validate_coord_in_dims(coord, self.dimensions)

        nid = pack_coords(*coord)
        if nid in self.neurons:
            raise KeyError(f"Neuron at {coord} already exists (ID: {nid})")

        # Create neuron with config
        neuron = create_neuron(
            neuron_id=nid,
            neuron_type=neuron_type,
            config=self.neuron_config,
            **kwargs,
        )

        self.neurons[nid] = neuron

        def mark_neuron_dirty(neuron_id: int = nid) -> None:
            self._dirty_neuron_ids.add(neuron_id)

        neuron.set_dirty_callback(mark_neuron_dirty)
        self.synapses[nid] = []
        self.in_degree[nid] = 0
        self._dirty_neuron_ids.add(nid)

        return nid

    def remove_neuron(self, neuron_id: int) -> bool:
        """Remove a neuron and all its connections.

        Args:
            neuron_id: ID of the neuron to remove.

        Returns:
            True if the neuron was removed, False if it didn't exist.
        """
        if neuron_id not in self.neurons:
            return False

        # Remove incoming synapses from other neurons to this neuron
        for pre_id, syn_list in list(self.synapses.items()):
            if pre_id == neuron_id:
                continue
            kept: list[Synapse] = []
            for syn in syn_list:
                if syn.target_id == neuron_id:
                    self._synapse_count -= 1
                    self.in_degree[neuron_id] = max(
                        0, self.in_degree.get(neuron_id, 0) - 1
                    )
                else:
                    kept.append(syn)
            self.synapses[pre_id] = kept

        # Remove outgoing synapses from this neuron
        outgoing = self.synapses.pop(neuron_id, [])
        for syn in outgoing:
            if syn.target_id in self.in_degree:
                self.in_degree[syn.target_id] = max(
                    0, self.in_degree[syn.target_id] - 1
                )
            self._synapse_count -= 1

        # Remove the neuron
        del self.neurons[neuron_id]
        self.in_degree.pop(neuron_id, None)
        self.input_cells.discard(neuron_id)
        self.output_cells.discard(neuron_id)
        self._dirty_neuron_ids.add(neuron_id)

        return True

    def get_neuron(self, neuron_id: int) -> Neuron | None:
        """Get a neuron by ID.

        Args:
            neuron_id: ID of the neuron.

        Returns:
            The Neuron instance, or None if not found.
        """
        return self.neurons.get(neuron_id)

    def get_neuron_at_coord(self, coord: Coord5D) -> Neuron | None:
        """Get a neuron at a specific coordinate.

        Args:
            coord: 5D coordinate.

        Returns:
            The Neuron instance, or None if not found.
        """
        nid = pack_coords(*coord)
        return self.neurons.get(nid)

    def has_neuron(self, neuron_id: int) -> bool:
        """Check if a neuron exists.

        Args:
            neuron_id: ID of the neuron.

        Returns:
            True if the neuron exists.
        """
        return neuron_id in self.neurons

    @property
    def neuron_count(self) -> int:
        """Number of neurons in the network."""
        return len(self.neurons)

    def neuron_ids(self) -> set[int]:
        """Get all neuron IDs."""
        return set(self.neurons.keys())

    # ========================================================================
    # Synapse Management
    # ========================================================================

    def connect(
        self,
        pre_id: int,
        post_id: int,
        weight: float,
        delay: int,
        config: SynapseConfig | None = None,
    ) -> bool:
        """Create a synaptic connection between two neurons.

        Args:
            pre_id: ID of the presynaptic neuron.
            post_id: ID of the postsynaptic neuron.
            weight: Synaptic weight (connection strength).
            delay: Transmission delay in ticks (1 - max_delay).
            config: Optional custom synapse configuration.

        Returns:
            True if the connection was created, False otherwise.

        Raises:
            ValueError: If neurons don't exist, delay is invalid,
                or connection rules are violated.
        """
        if pre_id not in self.neurons:
            raise ValueError(f"Presynaptic neuron {pre_id} not found")
        if post_id not in self.neurons:
            raise ValueError(f"Postsynaptic neuron {post_id} not found")
        if delay < 1 or delay > self.max_delay:
            raise ValueError(f"Delay must be 1..{self.max_delay}")
        if pre_id == post_id and not self.topology_config.allow_self_connections:
            raise ValueError("Self-connections are disabled")
        if not self.topology_config.allow_parallel_connections:
            if any(s.target_id == post_id for s in self.synapses[pre_id]):
                raise ValueError("Parallel connection already exists")

        # Create synapse
        synapse = create_synapse(post_id, weight, delay, config or self.synapse_config)
        self.synapses[pre_id].append(synapse)

        def mark_synapse_dirty(
            source_id: int = pre_id,
            target_id: int = post_id,
        ) -> None:
            self._dirty_synapse_ids.add((source_id, target_id))

        synapse.set_dirty_callback(mark_synapse_dirty)
        self._synapse_count += 1
        self.in_degree[post_id] = self.in_degree.get(post_id, 0) + 1
        self._dirty_synapse_ids.add((pre_id, post_id))

        return True

    def disconnect(self, pre_id: int, post_id: int) -> bool:
        """Remove a synaptic connection.

        Args:
            pre_id: ID of the presynaptic neuron.
            post_id: ID of the postsynaptic neuron.

        Returns:
            True if the connection was removed, False if it didn't exist.
        """
        if pre_id not in self.synapses:
            return False

        old_len = len(self.synapses[pre_id])
        self.synapses[pre_id] = [
            s for s in self.synapses[pre_id] if s.target_id != post_id
        ]
        removed = old_len - len(self.synapses[pre_id])

        if removed > 0:
            self._synapse_count -= removed
            if post_id in self.in_degree:
                self.in_degree[post_id] = max(0, self.in_degree[post_id] - removed)
            self._dirty_synapse_ids.add((pre_id, post_id))
            return True

        return False

    def get_synapses(self, pre_id: int) -> list[Synapse]:
        """Get all synapses from a neuron.

        Args:
            pre_id: ID of the presynaptic neuron.

        Returns:
            List of Synapse objects.
        """
        return self.synapses.get(pre_id, [])

    def get_incoming_synapses(self, post_id: int) -> list[tuple[int, Synapse]]:
        """Get all synapses targeting a neuron.

        Args:
            post_id: ID of the postsynaptic neuron.

        Returns:
            List of (presynaptic_neuron_id, Synapse) tuples.
        """
        incoming: list[tuple[int, Synapse]] = []
        for pre_id, syn_list in self.synapses.items():
            for syn in syn_list:
                if syn.target_id == post_id:
                    incoming.append((pre_id, syn))
        return incoming

    @property
    def synapse_count(self) -> int:
        """Number of synapses in the network."""
        return self._synapse_count

    @property
    def queued_event_count(self) -> int:
        """Number of events currently queued in the event buffer."""
        return self._queued_event_count

    # ========================================================================
    # Connection Initialization
    # ========================================================================

    def initialize_random_connections(
        self,
        connections_per_neuron: int | None = None,
        radius: float | None = None,
        weight_range: tuple[float, float] | None = None,
    ) -> None:
        """Initialize random connections between neurons within a radius.

        Args:
            connections_per_neuron: Target connections per neuron.
                If None, uses config value.
            radius: Neighbour radius in 5D space.
                If None, uses config value.
            weight_range: (min, max) weight range.
                If None, uses config values.
        """
        if connections_per_neuron is None:
            connections_per_neuron = self.network_config.initial_connections_per_neuron
        if radius is None:
            radius = self.network_config.neighbour_radius
        if weight_range is None:
            weight_range = (
                self.network_config.weight_min,
                self.network_config.weight_max,
            )

        wmin, wmax = weight_range

        for pre_id in list(self.neurons.keys()):
            pre_coord = unpack_coords(pre_id)

            # Collect candidate neurons within radius
            candidates: list[int] = []
            for ncoord in iter_neighbour_coords(pre_coord, self.dimensions, radius):
                nid = pack_coords(*ncoord)
                if nid not in self.neurons:
                    continue
                if nid == pre_id and not self.topology_config.allow_self_connections:
                    continue
                candidates.append(nid)

            if not candidates:
                continue

            # Select random targets
            sample_size = min(connections_per_neuron, len(candidates))
            for post_id in self.rng.sample(candidates, sample_size):
                weight = self.rng.uniform(wmin, wmax)
                delay = self.rng.randint(1, self.max_delay)
                try:
                    self.connect(pre_id, post_id, weight, delay)
                except ValueError:
                    continue  # Skip invalid connections

    def connect_neighbours(
        self,
        radius: float,
        weight_range: tuple[float, float] = (0.0, 0.5),
        probability: float = 0.1,
    ) -> None:
        """Connect neurons probabilistically within a radius.

        Args:
            radius: Neighbour radius in 5D space.
            weight_range: (min, max) weight range.
            probability: Connection probability between neighbours.
        """
        wmin, wmax = weight_range

        for pre_id in list(self.neurons.keys()):
            pre_coord = unpack_coords(pre_id)

            for ncoord in iter_neighbour_coords(pre_coord, self.dimensions, radius):
                post_id = pack_coords(*ncoord)
                if post_id not in self.neurons:
                    continue
                if (
                    post_id == pre_id
                    and not self.topology_config.allow_self_connections
                ):
                    continue

                if self.rng.random() < probability:
                    weight = self.rng.uniform(wmin, wmax)
                    delay = self.rng.randint(1, self.max_delay)
                    try:
                        self.connect(pre_id, post_id, weight, delay)
                    except ValueError:
                        continue

    # ========================================================================
    # Current Injection
    # ========================================================================

    def inject_current(self, neuron_id: int, current: float) -> None:
        """Inject an external current into a neuron.

        Args:
            neuron_id: ID of the target neuron.
            current: Current value to inject (can be positive or negative).
        """
        if neuron_id in self.neurons:
            self.pending_currents[neuron_id] = (
                self.pending_currents.get(neuron_id, 0.0) + current
            )

    def inject_current_batch(self, currents: dict[int, float]) -> None:
        """Inject currents into multiple neurons.

        Args:
            currents: Dictionary mapping neuron_id -> current value.
        """
        for nid, current in currents.items():
            self.inject_current(nid, current)

    def clear_pending_currents(self) -> None:
        """Clear all pending currents."""
        self.pending_currents.clear()

    # ========================================================================
    # Input/Output Layer Management
    # ========================================================================

    def set_input_output_cells(
        self,
        input_dim: str,
        input_coord: int,
        output_dim: str,
        output_coord: int,
    ) -> None:
        """Set input and output cells based on dimension boundaries.

        Args:
            input_dim: Name of the input dimension ('x', 'y', 'z', 'd4', 'd5').
            input_coord: Coordinate value on the input dimension.
            output_dim: Name of the output dimension.
            output_coord: Coordinate value on the output dimension.

        Raises:
            ValueError: If dimension names are unknown.
        """
        if input_dim not in DIM_NAMES:
            raise ValueError(f"Unknown dimension: {input_dim}")
        if output_dim not in DIM_NAMES:
            raise ValueError(f"Unknown dimension: {output_dim}")

        self.input_cells.clear()
        self.output_cells.clear()

        input_idx = DIM_NAMES[input_dim]
        output_idx = DIM_NAMES[output_dim]

        for nid in self.neurons:
            coord = unpack_coords(nid)
            if coord[input_idx] == input_coord:
                self.input_cells.add(nid)
            if coord[output_idx] == output_coord:
                self.output_cells.add(nid)

    def is_input_cell(self, neuron_id: int) -> bool:
        """Check if a neuron is an input cell."""
        return neuron_id in self.input_cells

    def is_output_cell(self, neuron_id: int) -> bool:
        """Check if a neuron is an output cell."""
        return neuron_id in self.output_cells

    # ========================================================================
    # Simulation Step
    # ========================================================================

    def step(self) -> StepResult:
        """Execute one simulation tick.

        Returns:
            StepResult containing the results of this tick.

        Raises:
            RuntimeError: If queue invariants are violated (with debug mode).
        """
        start = time.perf_counter()
        tick = self.current_tick
        slot_index = tick % len(self.event_slots)

        # 1. Apply external currents
        external_currents = self.pending_currents.copy()
        self.pending_currents.clear()

        # 2. Deliver queued spike events in deterministic order
        synaptic_currents: dict[int, float] = {}
        events = sorted(
            self.event_slots[slot_index],
            key=lambda e: (e.delivery_tick, e.source_id, e.target_id),
        )

        for ev in events:
            if self.debug_invariants and ev.delivery_tick != tick:
                raise RuntimeError(
                    f"Queue invariant violated: tick={tick}, delivery={ev.delivery_tick}"
                )
            if ev.target_id in self.neurons:
                synaptic_currents[ev.target_id] = (
                    synaptic_currents.get(ev.target_id, 0.0) + ev.weight
                )
            self.total_events_processed += 1

        delivered = len(events)
        self.event_slots[slot_index] = []
        self._queued_event_count -= delivered

        if self.debug_invariants and self._queued_event_count < 0:
            raise RuntimeError("queued_event_count became negative")

        # 3. Update neurons
        spike_ids: list[int] = []
        output_spikes: list[int] = []
        neuron_activity: dict[int, bool] = {}

        active = len(self.neurons)
        sum_v = 0.0
        sum_energy = 0.0
        min_v = float("inf")
        max_v = -float("inf")

        # Explicit deterministic iteration: sort by neuron_id
        # This ensures identical tick execution regardless of dict insertion order.
        for nid, neuron in sorted(self.neurons.items()):
            ext = external_currents.get(nid, 0.0)
            syn = synaptic_currents.get(nid, 0.0)
            neuron.last_external_current = ext
            neuron.last_synaptic_current = syn

            spiked = neuron.step(ext + syn, tick)
            neuron_activity[nid] = spiked

            sum_v += neuron.v
            sum_energy += neuron.energy
            min_v = min(min_v, neuron.v)
            max_v = max(max_v, neuron.v)

            if spiked:
                spike_ids.append(nid)
                self.total_spikes += 1

                if nid in self.output_cells:
                    output_spikes.append(nid)

                # Queue outgoing spikes in deterministic order (by target_id)
                for connection in sorted(
                    self.synapses.get(nid, []),
                    key=lambda s: s.target_id,
                ):
                    connection.last_pre_spike = tick
                    connection.mark_dirty()
                    delivery_tick = tick + connection.delay
                    slot = delivery_tick % len(self.event_slots)
                    self.event_slots[slot].append(
                        SpikeEvent(
                            nid,
                            connection.target_id,
                            connection.weight,
                            delivery_tick,
                        )
                    )
                    self._queued_event_count += 1

        # 4. Compute statistics
        if active:
            mean_v = sum_v / active
            mean_energy = sum_energy / active
        else:
            mean_v = min_v = max_v = mean_energy = 0.0

        # 5. Advance tick
        self.current_tick = tick + 1
        self._step_count += 1

        # 6. Debug invariants
        if self.debug_invariants:
            actual = sum(len(s) for s in self.event_slots)
            if actual != self._queued_event_count:
                raise RuntimeError(
                    f"Queue accounting mismatch: counter={self._queued_event_count}, actual={actual}"
                )

        elapsed = (time.perf_counter() - start) * 1000.0

        # 7. Build result
        result = StepResult(
            tick=tick,
            spike_ids=tuple(spike_ids),
            output_spike_ids=tuple(output_spikes),
            spikes_this_tick=len(spike_ids),
            total_spikes=self.total_spikes,
            delivered_events=delivered,
            queued_events=self._queued_event_count,
            external_injection_count=len(external_currents),
            external_total_current=sum(external_currents.values()),
            synaptic_current_targets=len(synaptic_currents),
            mean_v=mean_v,
            min_v=min_v,
            max_v=max_v,
            mean_energy=mean_energy,
            core_step_ms=elapsed,
            neuron_activity=neuron_activity,
            total_synapses=self._synapse_count,
            dirty_neuron_ids=tuple(sorted(self._dirty_neuron_ids)),
            dirty_synapse_ids=tuple(sorted(self._dirty_synapse_ids)),
        )

        # 8. Run post-step hooks
        for hook in tuple(self._post_step_hooks):
            try:
                hook(result)
            except Exception:
                _LOGGER.exception("Post-step hook failed at tick %s", result.tick)

        return result

    def step_batch(self, count: int) -> tuple[StepResult, ...]:
        """Execute consecutive ticks using the same semantics as ``step``.

        The batch is deliberately a thin native loop: every tick still drains
        its own event slot, advances the clock once, and runs post-step hooks.
        This makes it suitable for an equivalence check against repeated
        single-tick execution without changing the deterministic state model.
        """
        if isinstance(count, bool):
            raise TypeError("count must be an integer")
        if count < 1:
            raise ValueError("count must be >= 1")
        return tuple(self.step() for _ in range(count))

    # ========================================================================
    # Post-Step Hooks
    # ========================================================================

    def add_post_step_hook(self, hook: PostStepHook) -> None:
        """Register a hook that runs after each tick.

        Args:
            hook: Callback function receiving the StepResult.
        """
        if hook not in self._post_step_hooks:
            self._post_step_hooks.append(hook)

    def remove_post_step_hook(self, hook: PostStepHook) -> None:
        """Remove a previously registered hook."""
        try:
            self._post_step_hooks.remove(hook)
        except ValueError:
            pass

    def clear_post_step_hooks(self) -> None:
        """Remove all post-step hooks."""
        self._post_step_hooks.clear()

    # ========================================================================
    # State Inspection
    # ========================================================================

    def get_state_summary(self) -> dict[str, Any]:
        """Get a summary of the network state.

        Returns:
            Dictionary with network statistics.
        """
        return {
            "tick": self.current_tick,
            "neurons": len(self.neurons),
            "synapses": self._synapse_count,
            "input_cells": len(self.input_cells),
            "output_cells": len(self.output_cells),
            "total_spikes": self.total_spikes,
            "queued_events": self._queued_event_count,
            "step_count": self._step_count,
            "dimensions": self.dimensions,
        }

    def get_neurons_by_type(self) -> dict[NeuronType, list[int]]:
        """Get neurons grouped by type.

        Returns:
            Dictionary mapping NeuronType to list of neuron IDs.
        """
        result: dict[NeuronType, list[int]] = {}
        for nid, neuron in self.neurons.items():
            neuron_type = neuron.neuron_type
            result.setdefault(neuron_type, []).append(nid)
        return result

    def get_activity_metrics(self) -> dict[str, Any]:
        """Get activity metrics for the network.

        Returns:
            Dictionary with activity statistics.
        """
        if not self.neurons:
            return {"active_neurons": 0, "mean_firing_rate": 0.0}

        rates = [n.firing_rate_estimate for n in self.neurons.values()]
        return {
            "active_neurons": sum(1 for r in rates if r > 0.1),
            "mean_firing_rate": sum(rates) / len(rates),
            "max_firing_rate": max(rates) if rates else 0.0,
            "min_firing_rate": min(rates) if rates else 0.0,
        }

    def get_energy_stats(self) -> dict[str, Any]:
        """Get energy statistics for the network.

        Returns:
            Dictionary with energy statistics.
        """
        if not self.neurons:
            return {"mean_energy": 0.0, "min_energy": 0.0, "max_energy": 0.0}

        energies = [n.energy for n in self.neurons.values()]
        return {
            "mean_energy": sum(energies) / len(energies),
            "min_energy": min(energies),
            "max_energy": max(energies),
        }

    # ========================================================================
    # Serialization
    # ========================================================================

    def to_dict(self, include_state: bool = True) -> dict[str, Any]:
        """Serialize the network to a dictionary.

        Args:
            include_state: Whether to include neuron and synapse state.

        Returns:
            Dictionary containing network data.
        """
        data: dict[str, Any] = {
            "dimensions": list(self.dimensions),
            "tick": self.current_tick,
            "total_spikes": self.total_spikes,
            "synapse_count": self._synapse_count,
            "neurons": {},
            "synapses": {},
            "input_cells": list(self.input_cells),
            "output_cells": list(self.output_cells),
        }

        if include_state:
            # Serialize neurons
            for nid, neuron in self.neurons.items():
                data["neurons"][str(nid)] = neuron.to_dict()

            # Serialize synapses
            for pre_id, syn_list in self.synapses.items():
                data["synapses"][str(pre_id)] = [s.to_dict() for s in syn_list]

        return data

    @classmethod
    def from_dict(
        cls,
        data: dict[str, Any],
        rng: random.Random | None = None,
    ) -> NeuralNetwork:
        """Deserialize a network from a dictionary.

        Args:
            data: Dictionary containing network data.
            rng: Optional random generator.

        Returns:
            A new NeuralNetwork instance.

        Raises:
            ValueError: If the data is invalid.
        """
        # Create config from dimensions
        dims = tuple(data["dimensions"])
        config = Brain5DConfig(dimensions=dims)

        network = cls(config, rng)

        # Restore neurons
        for nid_str, neuron_data in data.get("neurons", {}).items():
            nid = int(nid_str)
            neuron = Neuron.from_dict(neuron_data)
            network.neurons[nid] = neuron
            network.synapses[nid] = []

        # Restore synapses
        for pre_id_str, syn_list in data.get("synapses", {}).items():
            pre_id = int(pre_id_str)
            for syn_data in syn_list:
                synapse = Synapse.from_dict(syn_data)
                network.synapses[pre_id].append(synapse)
                network._synapse_count += 1
                network.in_degree[synapse.target_id] = (
                    network.in_degree.get(synapse.target_id, 0) + 1
                )

        # Restore input/output cells
        network.input_cells = set(data.get("input_cells", []))
        network.output_cells = set(data.get("output_cells", []))

        # Restore state
        network.current_tick = data.get("tick", 0)
        network.total_spikes = data.get("total_spikes", 0)

        return network

    # ========================================================================
    # String Representation
    # ========================================================================

    def __str__(self) -> str:
        return (
            f"NeuralNetwork(neurons={len(self.neurons)}, "
            f"synapses={self._synapse_count}, "
            f"tick={self.current_tick}, "
            f"input={len(self.input_cells)}, output={len(self.output_cells)})"
        )

    def __repr__(self) -> str:
        return self.__str__()


# ============================================================================
# Factory Functions
# ============================================================================


def create_network(
    dimensions: Dim5D = (50, 50, 50, 50, 50),
    seed: int | None = None,
    **kwargs: Any,
) -> NeuralNetwork:
    """Create a neural network with default configuration.

    Args:
        dimensions: 5D dimensions for the network.
        seed: Optional seed for random number generator.
        **kwargs: Additional configuration parameters.

    Returns:
        A new NeuralNetwork instance.

    Example:
        >>> network = create_network((10, 10, 10, 10, 10), seed=42)
        >>> network.add_neuron((1, 2, 3, 4, 5))
    """
    rng = random.Random(seed) if seed is not None else random.Random()
    config = Brain5DConfig(dimensions=dimensions, **kwargs)
    return NeuralNetwork(config, rng)


# ============================================================================
# Module Exports
# ============================================================================

__all__ = [
    # Configuration
    "SimulationConfig",
    "TopologyConfig",
    "NetworkConfig",
    "Brain5DConfig",
    # Events
    "SpikeEvent",
    "StepResult",
    # Main class
    "NeuralNetwork",
    # Factory
    "create_network",
    # Types
    "PostStepHook",
    "ConfigDict",
]

# Preferred public name; preserve class identity for old configurations.
MHRNConfig = Brain5DConfig