src/storage/checkpoint.py

"""Typed runtime checkpoint sidecar for deterministic restore-and-continue."""

from __future__ import annotations

import json
import random
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol, cast, runtime_checkable


class CheckpointEventLike(Protocol):
    """Queued spike event fields required by the checkpoint sidecar."""

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


class CheckpointNetworkLike(Protocol):
    """Runtime fields required for the base checkpoint contract."""

    rng: random.Random
    current_tick: int
    total_spikes: int
    total_events_processed: int
    pending_currents: Mapping[int, float]
    input_cells: set[int]
    output_cells: set[int]
    event_slots: Sequence[Sequence[CheckpointEventLike]]


class CheckpointNeuronLike(Protocol):
    """Exact neuron fields required for deterministic continuation."""

    neuron_id: int
    a: float
    b: float
    c: float
    d: float
    v: float
    u: float
    energy: float
    spike_cost: float
    spike_counter: int
    last_spike_tick: int
    threshold_adaptation: float
    last_external_current: float
    last_synaptic_current: float
    firing_rate_estimate: float
    pre_trace: float
    post_trace: float
    _spike_count_window: int
    _last_update_tick: int


class CheckpointSynapseLike(Protocol):
    """Exact synapse fields required for deterministic continuation."""

    target_id: int
    weight: float
    delay: int
    eligibility: float
    last_pre_spike: int


@runtime_checkable
class ExactCheckpointNetworkLike(Protocol):
    """Optional exact network surface available on the real core."""

    neurons: Mapping[int, CheckpointNeuronLike]
    synapses: Mapping[int, Sequence[CheckpointSynapseLike]]


@dataclass(frozen=True, slots=True)
class RandomStateRecord:
    """JSON-safe representation of ``random.Random`` state."""

    version: int
    state: tuple[int, ...]
    gauss_next: float | None


@dataclass(frozen=True, slots=True)
class QueuedEventRecord:
    """One future spike delivery retained across restart."""

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


@dataclass(frozen=True, slots=True)
class NeuronRuntimeRecord:
    """Exact neuron state overlaid after compact snapshot recovery."""

    neuron_id: int
    a: float
    b: float
    c: float
    d: float
    membrane_v: float
    recovery_u: float
    energy: float
    spike_cost: float
    threshold_adaptation: float
    spike_counter: int
    last_spike_tick: int
    last_external_current: float
    last_synaptic_current: float
    # Internal neuron state required for deterministic continuation
    firing_rate_estimate: float = 0.0
    spike_count_window: int = 0
    last_update_tick: int = 0
    pre_trace: float = 0.0
    post_trace: float = 0.0


@dataclass(frozen=True, slots=True)
class SynapseRuntimeRecord:
    """Exact synapse state overlaid after compact snapshot recovery."""

    source_id: int
    target_id: int
    weight: float
    delay: int
    eligibility: float
    last_pre_spike: int


@dataclass(frozen=True, slots=True)
class HomeostasisRuntimeRecord:
    """Homeostasis engine state for deterministic continuation."""

    neuron_id: int
    rate_hz: float


@dataclass(frozen=True, slots=True)
class PendingRewardRecord:
    """A reward signal pending future application."""

    value: float
    tick: int


@dataclass(frozen=True, slots=True)
class LearningRuntimeRecord:
    """Learning engine state for deterministic continuation.

    Includes per-synapse traces and pending reward signals that
    affect future synaptic weight changes.
    """

    pre_id: int
    target_id: int
    last_pre_tick: int | None
    last_post_tick: int | None
    eligibility_value: float
    eligibility_last_tick: int | None


@dataclass(frozen=True, slots=True)
class RuntimeCheckpoint:
    """Runtime state required for deterministic restore-and-continue.

    Version history:
        1: Initial checkpoint format
        2: Added neuron_states and synapse_states
        3: Added pending_currents, input_cells, output_cells
        4: Added homeostasis_state and learning_state
    """

    version: int
    current_tick: int
    total_spikes: int
    total_events_processed: int
    rng: RandomStateRecord
    pending_currents: tuple[tuple[int, float], ...]
    input_cells: tuple[int, ...]
    output_cells: tuple[int, ...]
    queued_events: tuple[QueuedEventRecord, ...]
    neuron_states: tuple[NeuronRuntimeRecord, ...] = ()
    synapse_states: tuple[SynapseRuntimeRecord, ...] = ()
    homeostasis_state: tuple[HomeostasisRuntimeRecord, ...] = ()
    learning_state: tuple[LearningRuntimeRecord, ...] = ()
    pending_rewards: tuple[PendingRewardRecord, ...] = ()


def capture_runtime_checkpoint(
    network: CheckpointNetworkLike,
    *,
    homeostasis_rates: dict[int, float] | None = None,
    learning_states: list[dict[str, object]] | None = None,
    pending_rewards: list[dict[str, object]] | None = None,
) -> RuntimeCheckpoint:
    """Capture runtime state and exact core state when available.

    Args:
        network: The network to capture state from.
        homeostasis_rates: Optional dict of neuron_id -> smoothed firing rate
            from the homeostasis engine.
        learning_states: Optional list of learning engine state dicts,
            each containing pre_id, target_id, last_pre_tick, last_post_tick,
            and eligibility_value.
        pending_rewards: Optional list of pending reward signal dicts,
            each containing value and tick.

    Returns:
        A RuntimeCheckpoint with version 4 (includes homeostasis + learning).
    """
    version, raw_state, gauss_next = network.rng.getstate()
    queued = tuple(
        QueuedEventRecord(
            source_id=int(event.source_id),
            target_id=int(event.target_id),
            weight=float(event.weight),
            delivery_tick=int(event.delivery_tick),
        )
        for slot in network.event_slots
        for event in slot
    )
    neuron_states: tuple[NeuronRuntimeRecord, ...] = ()
    synapse_states: tuple[SynapseRuntimeRecord, ...] = ()
    if isinstance(network, ExactCheckpointNetworkLike):
        neuron_states = tuple(
            NeuronRuntimeRecord(
                neuron_id=int(neuron_id),
                a=float(neuron.a),
                b=float(neuron.b),
                c=float(neuron.c),
                d=float(neuron.d),
                membrane_v=float(neuron.v),
                recovery_u=float(neuron.u),
                energy=float(neuron.energy),
                spike_cost=float(neuron.spike_cost),
                threshold_adaptation=float(neuron.threshold_adaptation),
                spike_counter=int(neuron.spike_counter),
                last_spike_tick=int(neuron.last_spike_tick),
                last_external_current=float(neuron.last_external_current),
                last_synaptic_current=float(neuron.last_synaptic_current),
                firing_rate_estimate=float(neuron.firing_rate_estimate),
                spike_count_window=int(
                    neuron._spike_count_window  # pyright: ignore[reportPrivateUsage]
                ),
                last_update_tick=int(
                    neuron._last_update_tick  # pyright: ignore[reportPrivateUsage]
                ),
                pre_trace=float(neuron.pre_trace),
                post_trace=float(neuron.post_trace),
            )
            for neuron_id, neuron in sorted(network.neurons.items())
        )
        synapse_states = tuple(
            SynapseRuntimeRecord(
                source_id=int(source_id),
                target_id=int(synapse.target_id),
                weight=float(synapse.weight),
                delay=int(synapse.delay),
                eligibility=float(synapse.eligibility),
                last_pre_spike=int(synapse.last_pre_spike),
            )
            for source_id, synapses in sorted(network.synapses.items())
            for synapse in synapses
        )

    # Capture homeostasis state
    homeostasis_state: tuple[HomeostasisRuntimeRecord, ...] = ()
    if homeostasis_rates is not None:
        homeostasis_state = tuple(
            HomeostasisRuntimeRecord(neuron_id=int(nid), rate_hz=float(rate))
            for nid, rate in sorted(homeostasis_rates.items())
        )

    # Capture learning state
    learning_state: tuple[LearningRuntimeRecord, ...] = ()
    if learning_states is not None:
        learning_state = tuple(
            LearningRuntimeRecord(
                pre_id=int(cast("int | str", s["pre_id"])),
                target_id=int(cast("int | str", s["target_id"])),
                last_pre_tick=(
                    int(cast("int | str", s["last_pre_tick"]))
                    if s.get("last_pre_tick") is not None
                    else None
                ),
                last_post_tick=(
                    int(cast("int | str", s["last_post_tick"]))
                    if s.get("last_post_tick") is not None
                    else None
                ),
                eligibility_value=float(
                    cast("int | float | str", s.get("eligibility_value", 0.0))
                ),
                eligibility_last_tick=(
                    int(cast("int | str", s["eligibility_last_tick"]))
                    if s.get("eligibility_last_tick") is not None
                    else None
                ),
            )
            for s in learning_states
        )

    # Capture pending rewards
    pending_reward_records: tuple[PendingRewardRecord, ...] = ()
    if pending_rewards is not None:
        pending_reward_records = tuple(
            PendingRewardRecord(
                value=float(cast("int | float | str", r["value"])),
                tick=int(cast("int | str", r["tick"])),
            )
            for r in pending_rewards
        )

    return RuntimeCheckpoint(
        version=4,
        current_tick=int(network.current_tick),
        total_spikes=int(network.total_spikes),
        total_events_processed=int(network.total_events_processed),
        rng=RandomStateRecord(
            int(version), tuple(int(value) for value in raw_state), gauss_next
        ),
        pending_currents=tuple(
            sorted(
                (int(key), float(value))
                for key, value in network.pending_currents.items()
            )
        ),
        input_cells=tuple(sorted(int(value) for value in network.input_cells)),
        output_cells=tuple(sorted(int(value) for value in network.output_cells)),
        queued_events=queued,
        neuron_states=neuron_states,
        synapse_states=synapse_states,
        homeostasis_state=homeostasis_state,
        learning_state=learning_state,
        pending_rewards=pending_reward_records,
    )


def write_runtime_checkpoint(path: Path, checkpoint: RuntimeCheckpoint) -> None:
    """Write a deterministic JSON runtime sidecar."""
    payload: dict[str, object] = {
        "version": checkpoint.version,
        "current_tick": checkpoint.current_tick,
        "total_spikes": checkpoint.total_spikes,
        "total_events_processed": checkpoint.total_events_processed,
        "rng": {
            "version": checkpoint.rng.version,
            "state": list(checkpoint.rng.state),
            "gauss_next": checkpoint.rng.gauss_next,
        },
        "pending_currents": [list(item) for item in checkpoint.pending_currents],
        "input_cells": list(checkpoint.input_cells),
        "output_cells": list(checkpoint.output_cells),
        "queued_events": [
            {
                "source_id": event.source_id,
                "target_id": event.target_id,
                "weight": event.weight,
                "delivery_tick": event.delivery_tick,
            }
            for event in checkpoint.queued_events
        ],
        "neuron_states": [
            {
                "neuron_id": state.neuron_id,
                "a": state.a,
                "b": state.b,
                "c": state.c,
                "d": state.d,
                "membrane_v": state.membrane_v,
                "recovery_u": state.recovery_u,
                "energy": state.energy,
                "spike_cost": state.spike_cost,
                "threshold_adaptation": state.threshold_adaptation,
                "spike_counter": state.spike_counter,
                "last_spike_tick": state.last_spike_tick,
                "last_external_current": state.last_external_current,
                "last_synaptic_current": state.last_synaptic_current,
                "firing_rate_estimate": state.firing_rate_estimate,
                "spike_count_window": state.spike_count_window,
                "last_update_tick": state.last_update_tick,
                "pre_trace": state.pre_trace,
                "post_trace": state.post_trace,
            }
            for state in checkpoint.neuron_states
        ],
        "synapse_states": [
            {
                "source_id": state.source_id,
                "target_id": state.target_id,
                "weight": state.weight,
                "delay": state.delay,
                "eligibility": state.eligibility,
                "last_pre_spike": state.last_pre_spike,
            }
            for state in checkpoint.synapse_states
        ],
        "homeostasis_state": [
            {
                "neuron_id": state.neuron_id,
                "rate_hz": state.rate_hz,
            }
            for state in checkpoint.homeostasis_state
        ],
        "learning_state": [
            {
                "pre_id": state.pre_id,
                "target_id": state.target_id,
                "last_pre_tick": state.last_pre_tick,
                "last_post_tick": state.last_post_tick,
                "eligibility_value": state.eligibility_value,
                "eligibility_last_tick": state.eligibility_last_tick,
            }
            for state in checkpoint.learning_state
        ],
        "pending_rewards": [
            {
                "value": r.value,
                "tick": r.tick,
            }
            for r in checkpoint.pending_rewards
        ],
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(payload, sort_keys=True, separators=(",", ":")),
        encoding="utf-8",
    )


def _mapping(value: object, name: str) -> dict[str, object]:
    """Validate and cast a JSON object to dict[str, object]."""
    if not isinstance(value, dict):
        raise ValueError(f"{name} must be an object")
    # Cast to dict[str, object] after validation
    value = cast(dict[str, object], value)
    result: dict[str, object] = {}
    for key, item in value.items():
        # key is already str due to the cast, so no need for isinstance check
        result[key] = item
    return result


def _list(value: object, name: str) -> list[object]:
    """Validate and cast a JSON array to list[object]."""
    if not isinstance(value, list):
        raise ValueError(f"{name} must be a list")
    return cast(list[object], value)


def _int(value: object, name: str) -> int:
    """Validate and cast a JSON value to int."""
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{name} must be an integer")
    return value


def _float(value: object, name: str) -> float:
    """Validate and cast a JSON value to float."""
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{name} must be numeric")
    return float(value)


def _optional_float(value: object, name: str) -> float | None:
    """Validate and cast a JSON value to optional float."""
    if value is None:
        return None
    return _float(value, name)


def read_runtime_checkpoint(path: Path) -> RuntimeCheckpoint:
    """Read and strictly validate checkpoint versions 1 through 4."""
    loaded: object = json.loads(path.read_text(encoding="utf-8"))
    raw = _mapping(loaded, "checkpoint")
    version = _int(raw.get("version", 0), "version")
    if version not in {1, 2, 3, 4}:
        raise ValueError("unsupported runtime checkpoint")
    rng_raw = _mapping(raw.get("rng"), "rng")
    state_raw = _list(rng_raw.get("state"), "rng.state")

    pending: list[tuple[int, float]] = []
    for index, item in enumerate(_list(raw.get("pending_currents", []), "pending")):
        pair = _list(item, f"pending_currents[{index}]")
        if len(pair) != 2:
            raise ValueError("invalid pending-current record")
        pending.append(
            (
                _int(pair[0], f"pending_currents[{index}].id"),
                _float(pair[1], f"pending_currents[{index}].value"),
            )
        )

    events: list[QueuedEventRecord] = []
    for index, item in enumerate(_list(raw.get("queued_events", []), "events")):
        event = _mapping(item, f"queued_events[{index}]")
        events.append(
            QueuedEventRecord(
                source_id=_int(event.get("source_id"), "event.source_id"),
                target_id=_int(event.get("target_id"), "event.target_id"),
                weight=_float(event.get("weight"), "event.weight"),
                delivery_tick=_int(event.get("delivery_tick"), "event.delivery_tick"),
            )
        )

    neuron_states: list[NeuronRuntimeRecord] = []
    for index, item in enumerate(_list(raw.get("neuron_states", []), "neurons")):
        state = _mapping(item, f"neuron_states[{index}]")
        neuron_states.append(
            NeuronRuntimeRecord(
                neuron_id=_int(state.get("neuron_id"), "neuron.neuron_id"),
                a=_float(state.get("a"), "neuron.a"),
                b=_float(state.get("b"), "neuron.b"),
                c=_float(state.get("c"), "neuron.c"),
                d=_float(state.get("d"), "neuron.d"),
                membrane_v=_float(state.get("membrane_v"), "neuron.v"),
                recovery_u=_float(state.get("recovery_u"), "neuron.u"),
                energy=_float(state.get("energy"), "neuron.energy"),
                spike_cost=_float(state.get("spike_cost"), "neuron.spike_cost"),
                threshold_adaptation=_float(
                    state.get("threshold_adaptation"), "neuron.threshold"
                ),
                spike_counter=_int(state.get("spike_counter"), "neuron.spike_counter"),
                last_spike_tick=_int(
                    state.get("last_spike_tick"), "neuron.last_spike_tick"
                ),
                last_external_current=_float(
                    state.get("last_external_current", 0.0),
                    "neuron.last_external_current",
                ),
                last_synaptic_current=_float(
                    state.get("last_synaptic_current", 0.0),
                    "neuron.last_synaptic_current",
                ),
                firing_rate_estimate=_float(
                    state.get("firing_rate_estimate", 0.0),
                    "neuron.firing_rate_estimate",
                ),
                spike_count_window=_int(
                    state.get("spike_count_window", 0),
                    "neuron.spike_count_window",
                ),
                last_update_tick=_int(
                    state.get("last_update_tick", 0),
                    "neuron.last_update_tick",
                ),
                pre_trace=_float(
                    state.get("pre_trace", 0.0),
                    "neuron.pre_trace",
                ),
                post_trace=_float(
                    state.get("post_trace", 0.0),
                    "neuron.post_trace",
                ),
            )
        )

    synapse_states: list[SynapseRuntimeRecord] = []
    for index, item in enumerate(_list(raw.get("synapse_states", []), "synapses")):
        state = _mapping(item, f"synapse_states[{index}]")
        synapse_states.append(
            SynapseRuntimeRecord(
                source_id=_int(state.get("source_id"), "synapse.source_id"),
                target_id=_int(state.get("target_id"), "synapse.target_id"),
                weight=_float(state.get("weight"), "synapse.weight"),
                delay=_int(state.get("delay"), "synapse.delay"),
                eligibility=_float(state.get("eligibility"), "synapse.eligibility"),
                last_pre_spike=_int(
                    state.get("last_pre_spike"), "synapse.last_pre_spike"
                ),
            )
        )

    input_raw = _list(raw.get("input_cells", []), "input_cells")
    output_raw = _list(raw.get("output_cells", []), "output_cells")

    # Read homeostasis state (version 4+)
    homeostasis_state: list[HomeostasisRuntimeRecord] = []
    for index, item in enumerate(
        _list(raw.get("homeostasis_state", []), "homeostasis")
    ):
        state = _mapping(item, f"homeostasis_state[{index}]")
        homeostasis_state.append(
            HomeostasisRuntimeRecord(
                neuron_id=_int(state.get("neuron_id"), "homeostasis.neuron_id"),
                rate_hz=_float(state.get("rate_hz"), "homeostasis.rate_hz"),
            )
        )

    # Read learning state (version 4+)
    learning_state: list[LearningRuntimeRecord] = []
    for index, item in enumerate(_list(raw.get("learning_state", []), "learning")):
        state = _mapping(item, f"learning_state[{index}]")
        learning_state.append(
            LearningRuntimeRecord(
                pre_id=_int(state.get("pre_id"), "learning.pre_id"),
                target_id=_int(state.get("target_id"), "learning.target_id"),
                last_pre_tick=(
                    _int(state["last_pre_tick"], "learning.last_pre_tick")
                    if state.get("last_pre_tick") is not None
                    else None
                ),
                last_post_tick=(
                    _int(state["last_post_tick"], "learning.last_post_tick")
                    if state.get("last_post_tick") is not None
                    else None
                ),
                eligibility_value=_float(
                    state.get("eligibility_value", 0.0), "learning.eligibility_value"
                ),
                eligibility_last_tick=(
                    _int(
                        state["eligibility_last_tick"], "learning.eligibility_last_tick"
                    )
                    if state.get("eligibility_last_tick") is not None
                    else None
                ),
            )
        )

    # Read pending rewards (version 4+)
    pending_rewards: list[PendingRewardRecord] = []
    for index, item in enumerate(
        _list(raw.get("pending_rewards", []), "pending_rewards")
    ):
        r = _mapping(item, f"pending_rewards[{index}]")
        pending_rewards.append(
            PendingRewardRecord(
                value=_float(r.get("value", 0.0), "pending_reward.value"),
                tick=_int(r.get("tick", 0), "pending_reward.tick"),
            )
        )

    return RuntimeCheckpoint(
        version=version,
        current_tick=_int(raw.get("current_tick"), "current_tick"),
        total_spikes=_int(raw.get("total_spikes"), "total_spikes"),
        total_events_processed=_int(
            raw.get("total_events_processed"), "total_events_processed"
        ),
        rng=RandomStateRecord(
            version=_int(rng_raw.get("version"), "rng.version"),
            state=tuple(_int(value, "rng.state item") for value in state_raw),
            gauss_next=_optional_float(rng_raw.get("gauss_next"), "rng.gauss_next"),
        ),
        pending_currents=tuple(pending),
        input_cells=tuple(_int(value, "input cell") for value in input_raw),
        output_cells=tuple(_int(value, "output cell") for value in output_raw),
        queued_events=tuple(events),
        neuron_states=tuple(neuron_states),
        synapse_states=tuple(synapse_states),
        homeostasis_state=tuple(homeostasis_state),
        learning_state=tuple(learning_state),
        pending_rewards=tuple(pending_rewards),
    )