Blame for src/storage/core_restore.py

"""Restore a real Brain-5D core from snapshot, journal, and runtime sidecar."""

from __future__ import annotations

import random
from pathlib import Path
from typing import Any, Protocol, assert_never, cast

from src.config.loader import ConfigDict
from src.core.network import NeuralNetwork, SpikeEvent
from src.core.spatial_index import unpack_coords

from .b5d import B5DReader
from .checkpoint import RuntimeCheckpoint, read_runtime_checkpoint
from .recovery import RecoveryManager
from .structural_journal import StructuralChangeKind, StructuralChangeRecord
from .structural_recovery import StructuralRecoveryManager


class StructuralRestoreManipulator(Protocol):
    def create_neuron(self, coord: tuple[int, int, int, int, int]) -> int: ...
    def delete_neuron(self, neuron_id: int) -> None: ...
    def create_synapse(
        self, source_id: int, target_id: int, weight: float, delay: int
    ) -> None: ...
    def delete_synapse(self, source_id: int, target_id: int) -> None: ...


class _StructuralReplayTarget:
    def __init__(self, manipulator: StructuralRestoreManipulator) -> None:
        self._manipulator = manipulator

    def apply_structural_record(self, record: StructuralChangeRecord) -> bool:
        if record.kind is StructuralChangeKind.NEURON_ADD:
            if record.coord is None:
                return False
            neuron_id = self._manipulator.create_neuron(record.coord)
            return record.neuron_id is None or neuron_id == record.neuron_id
        if record.kind is StructuralChangeKind.NEURON_REMOVE:
            if record.neuron_id is None:
                return False
            self._manipulator.delete_neuron(record.neuron_id)
            return True
        if record.kind is StructuralChangeKind.SYNAPSE_ADD:
            if (
                record.source_id is None
                or record.target_id is None
                or record.weight is None
                or record.delay is None
            ):
                return False
            self._manipulator.create_synapse(
                record.source_id,
                record.target_id,
                record.weight,
                record.delay,
            )
            return True
        if record.kind is StructuralChangeKind.SYNAPSE_REMOVE:
            if record.source_id is None or record.target_id is None:
                return False
            self._manipulator.delete_synapse(record.source_id, record.target_id)
            return True
        assert_never(record.kind)


class RestoredNeuralNetwork(NeuralNetwork):
    """NeuralNetwork subclass exposing a narrow restoration hook."""

    def restore_queued_count(self, value: int) -> None:
        """Restore queue counter after event slots have been reconstructed."""
        self._queued_event_count = value


def restore_network(
    snapshot_path: Path,
    journal_path: Path,
    checkpoint_path: Path,
    config: ConfigDict,
    recovered_path: Path,
    *,
    structural_journal_path: Path | None = None,
) -> RestoredNeuralNetwork:
    """Recover persisted state and return a runnable deterministic network.

    Args:
        snapshot_path: Path to the base `.b5d` snapshot.
        journal_path: Path to the `.b5d.journal` delta journal.
        checkpoint_path: Path to the runtime checkpoint JSON.
        config: Network configuration (dict).
        recovered_path: Path to write the recovered snapshot (if needed).
        structural_journal_path: Optional structural journal path.

    Returns:
        A restored RestoredNeuralNetwork instance.

    Raises:
        RuntimeError: If recovery fails.
        ValueError: If the snapshot is not restart-capable or required fields are missing.
    """
    result = RecoveryManager(snapshot_path, journal_path).recover(recovered_path)
    if not result.success:
        raise RuntimeError(result.error or "snapshot recovery failed")

    checkpoint = read_runtime_checkpoint(checkpoint_path)
    rng = _restore_rng(checkpoint)

    # Cast ConfigDict to dict[str, Any] for the NeuralNetwork constructor
    network = RestoredNeuralNetwork(cast(dict[str, Any], config), rng)

    with B5DReader(recovered_path) as reader:
        for neuron_record in reader.iter_neurons():
            neuron_id = network.add_neuron(unpack_coords(neuron_record.neuron_id))
            neuron = network.neurons[neuron_id]

            # Ensure restart-capable snapshot was used
            if neuron_record.a is None or neuron_record.b is None:
                raise ValueError("restart-capable snapshot required")
            if neuron_record.c is None or neuron_record.d is None:
                raise ValueError("restart-capable snapshot required")
            if neuron_record.spike_cost is None:
                raise ValueError("restart-capable snapshot required")
            if neuron_record.spike_counter is None:
                raise ValueError("restart-capable snapshot required")
            if neuron_record.last_spike_tick is None:
                raise ValueError("restart-capable snapshot required")

            # Restore core neuron parameters
            neuron.a = neuron_record.a
            neuron.b = neuron_record.b
            neuron.c = neuron_record.c
            neuron.d = neuron_record.d
            neuron.spike_cost = neuron_record.spike_cost
            neuron.spike_counter = neuron_record.spike_counter
            neuron.last_spike_tick = neuron_record.last_spike_tick

            # Restore dynamic state from optical record
            neuron.v = neuron_record.optical.membrane_v
            neuron.u = neuron_record.optical.recovery_u
            neuron.energy = neuron_record.optical.energy
            neuron.threshold_adaptation = neuron_record.optical.threshold_adaptation

        for synapse_record in reader.iter_synapses():
            network.connect(
                synapse_record.source_id,
                synapse_record.target_id,
                synapse_record.weight,
                synapse_record.delay,
            )
            synapse = network.synapses[synapse_record.source_id][-1]
            synapse.eligibility = synapse_record.eligibility
            synapse.last_pre_spike = synapse_record.last_pre_spike

    # Replay structural plasticity journal if provided
    if structural_journal_path is not None:
        from src.manipulation.manipulator import Brain5DManipulator

        structural_target = _StructuralReplayTarget(Brain5DManipulator(network))
        StructuralRecoveryManager().replay(
            structural_target,
            structural_journal_path,
        )

    # Restore runtime state and exact neuron/synapse values from checkpoint
    _restore_runtime(network, checkpoint)
    _restore_exact_neuron_state(network, checkpoint)
    _restore_exact_synapse_state(network, checkpoint)

    return network


def _restore_rng(checkpoint: RuntimeCheckpoint) -> random.Random:
    """Restore the random number generator from checkpoint state."""
    rng = random.Random()
    rng.setstate(
        (
            checkpoint.rng.version,
            tuple(checkpoint.rng.state),
            checkpoint.rng.gauss_next,
        )
    )
    return rng


def _restore_runtime(
    network: RestoredNeuralNetwork,
    checkpoint: RuntimeCheckpoint,
) -> None:
    """Restore runtime state from checkpoint."""
    network.current_tick = checkpoint.current_tick
    network.total_spikes = checkpoint.total_spikes
    network.total_events_processed = checkpoint.total_events_processed
    network.pending_currents = dict(checkpoint.pending_currents)
    network.input_cells = set(checkpoint.input_cells)
    network.output_cells = set(checkpoint.output_cells)

    # Rebuild event queue
    network.event_slots = [[] for _ in range(network.max_delay + 1)]
    queued = 0
    for event in checkpoint.queued_events:
        slot = event.delivery_tick % len(network.event_slots)
        network.event_slots[slot].append(
            SpikeEvent(
                source_id=event.source_id,
                target_id=event.target_id,
                weight=event.weight,
                delivery_tick=event.delivery_tick,
            )
        )
        queued += 1
    network.restore_queued_count(queued)


def _restore_exact_neuron_state(
    network: RestoredNeuralNetwork,
    checkpoint: RuntimeCheckpoint,
) -> None:
    """Overlay exact values lost by the compact snapshot representation."""
    for state in checkpoint.neuron_states:
        neuron = network.neurons.get(state.neuron_id)
        if neuron is None:
            raise ValueError(
                f"runtime checkpoint references missing neuron {state.neuron_id}"
            )
        neuron.a = state.a
        neuron.b = state.b
        neuron.c = state.c
        neuron.d = state.d
        neuron.v = state.membrane_v
        neuron.u = state.recovery_u
        neuron.energy = state.energy
        neuron.spike_cost = state.spike_cost
        neuron.threshold_adaptation = state.threshold_adaptation
        neuron.spike_counter = state.spike_counter
        neuron.last_spike_tick = state.last_spike_tick
        neuron.last_external_current = state.last_external_current
        neuron.last_synaptic_current = state.last_synaptic_current
        neuron.firing_rate_estimate = state.firing_rate_estimate
        neuron._spike_count_window = (  # pyright: ignore[reportPrivateUsage]
            state.spike_count_window
        )
        neuron._last_update_tick = (  # pyright: ignore[reportPrivateUsage]
            state.last_update_tick
        )
        neuron.pre_trace = state.pre_trace
        neuron.post_trace = state.post_trace


def _restore_exact_synapse_state(
    network: RestoredNeuralNetwork,
    checkpoint: RuntimeCheckpoint,
) -> None:
    """Overlay exact synaptic float state lost by V1 float32 fields."""
    if not checkpoint.synapse_states:
        return

    for state in checkpoint.synapse_states:
        candidates = network.synapses.get(state.source_id)
        if candidates is None:
            raise ValueError(
                f"runtime checkpoint references missing source {state.source_id}"
            )
        synapse = next(
            (item for item in candidates if item.target_id == state.target_id),
            None,
        )
        if synapse is None:
            raise ValueError(
                "runtime checkpoint references missing synapse "
                f"{state.source_id}->{state.target_id}"
            )
        synapse.weight = state.weight
        synapse.delay = state.delay
        synapse.eligibility = state.eligibility
        synapse.last_pre_spike = state.last_pre_spike


def restore_homeostasis_state(
    homeostasis_engine: Any,
    checkpoint: RuntimeCheckpoint,
) -> None:
    """Restore homeostasis engine state from checkpoint v4.

    Args:
        homeostasis_engine: A HomeostasisEngine instance (must have _rates_hz dict).
        checkpoint: The RuntimeCheckpoint with homeostasis_state.

    Raises:
        ValueError: If a neuron referenced in the checkpoint no longer exists.
    """
    if not checkpoint.homeostasis_state:
        return

    rates = {}
    for record in checkpoint.homeostasis_state:
        rates[int(record.neuron_id)] = float(record.rate_hz)

    # Restore the internal rates dict
    if hasattr(homeostasis_engine, "_rates_hz"):
        homeostasis_engine._rates_hz = rates


def restore_learning_state(
    learning_engine: Any,
    checkpoint: RuntimeCheckpoint,
) -> None:
    """Restore learning engine state from checkpoint v4.

    Restores per-synapse traces (last_pre_tick, last_post_tick, eligibility)
    and pending reward signals.

    Args:
        learning_engine: A LearningEngine instance (must have _states dict
            and _pending_rewards list).
        checkpoint: The RuntimeCheckpoint with learning_state and pending_rewards.

    Raises:
        ValueError: If a synapse referenced in the checkpoint no longer exists.
    """
    if not hasattr(learning_engine, "_states"):
        return

    # Restore per-synapse traces
    # Uses stable (pre_id, target_id) key matching LearningEngine's
    # deterministic topology indexing.
    if checkpoint.learning_state:
        states = learning_engine._states
        for record in checkpoint.learning_state:
            pre_id = int(record.pre_id)
            target_id = int(record.target_id)
            key = (pre_id, target_id)
            if key in states:
                state = states[key]
                state.last_pre_tick = record.last_pre_tick
                state.last_post_tick = record.last_post_tick
                # EligibilityTrace uses .value, not ._trace
                state.eligibility.value = record.eligibility_value
                # Restore per-trace last_tick so decay is accurate
                state.eligibility.last_tick = record.eligibility_last_tick

    # Restore pending rewards (independent of learning_state presence)
    if checkpoint.pending_rewards and hasattr(learning_engine, "_pending_rewards"):
        from src.learning.reward import RewardSignal

        learning_engine._pending_rewards = [
            RewardSignal(value=float(r.value), tick=int(r.tick))
            for r in checkpoint.pending_rewards
        ]


# ============================================================================
# Production Restore Bundle
# ============================================================================


class RestoredBundle:
    """Complete production restore result: network + optional engines.

    This bundles the restored neural network together with its optional
    homeostasis and learning engines, providing a single restore-and-continue
    entry point for production use.

    Attributes:
        network: The restored RestoredNeuralNetwork instance.
        homeostasis_engine: Optional restored HomeostasisEngine, or None.
        learning_engine: Optional restored LearningEngine, or None.
    """

    def __init__(
        self,
        network: RestoredNeuralNetwork,
        homeostasis_engine: Any | None = None,
        learning_engine: Any | None = None,
    ) -> None:
        self.network = network
        self.homeostasis_engine = homeostasis_engine
        self.learning_engine = learning_engine


def restore_full(
    snapshot_path: Path,
    journal_path: Path,
    checkpoint_path: Path,
    config: ConfigDict,
    recovered_path: Path,
    *,
    structural_journal_path: Path | None = None,
    create_homeostasis_engine: bool = False,
    create_learning_engine: bool = False,
) -> RestoredBundle:
    """Restore network and optionally create and restore engines.

    This is the production entry point that replaces calling
    ``restore_network()`` followed by manual engine creation and restore.
    It ensures the full restore chain is executed atomically:

    1. Recover snapshot and restore network (same as restore_network())
    2. Optionally create HomeostasisEngine and restore its state
    3. Optionally create LearningEngine and restore its state

    Args:
        snapshot_path: Path to the base ``.b5d`` snapshot.
        journal_path: Path to the ``.b5d.journal`` delta journal.
        checkpoint_path: Path to the runtime checkpoint JSON.
        config: Network configuration (dict).
        recovered_path: Path to write the recovered snapshot.
        structural_journal_path: Optional structural journal path.
        create_homeostasis_engine: If True, create and restore
            HomeostasisEngine from checkpoint.
        create_learning_engine: If True, create and restore
            LearningEngine from checkpoint.

    Returns:
        A RestoredBundle containing the network and optionally the engines.

    Raises:
        RuntimeError: If recovery fails.
        ValueError: If the snapshot is not restart-capable.
    """
    network = restore_network(
        snapshot_path=snapshot_path,
        journal_path=journal_path,
        checkpoint_path=checkpoint_path,
        config=config,
        recovered_path=recovered_path,
        structural_journal_path=structural_journal_path,
    )

    checkpoint = read_runtime_checkpoint(checkpoint_path)
    homeostasis_engine: Any = None
    learning_engine: Any = None

    if create_homeostasis_engine:
        from src.homeostasis.engine import HomeostasisEngine

        homeostasis_engine = HomeostasisEngine(network, dict(config))
        restore_homeostasis_state(homeostasis_engine, checkpoint)
        homeostasis_engine.attach()

    if create_learning_engine:
        from src.learning.learning_engine import LearningEngine

        learning_engine = LearningEngine(network, dict(config))
        restore_learning_state(learning_engine, checkpoint)
        learning_engine.attach()

    return RestoredBundle(
        network=network,
        homeostasis_engine=homeostasis_engine,
        learning_engine=learning_engine,
    )