snnlab.sim.execution

class DenseArrayBinding:
    """One concrete dense tensor resolved against a named graph input."""

    input_id: str
    value: torch.Tensor
    source: Mapping[str, Any] = field(default_factory=dict)

class TargetArrayBinding:
    """One concrete target vector resolved against a named training target."""

    target_id: str
    value: torch.Tensor
    source: Mapping[str, Any] = field(default_factory=dict)

class EventStreamBinding:
    """Sparse binary spike events resolved against one named graph input."""

    input_id: str
    steps: torch.Tensor
    batches: torch.Tensor
    channels: torch.Tensor
    steps_count: int
    batch_size: int
    source: Mapping[str, Any] = field(default_factory=dict)

class PoissonInputBinding:
    """Generated Bernoulli-discretised Poisson spikes for one graph input."""

    input_id: str
    steps_count: int
    batch_size: int
    rates_hz: Sequence[float]
    seed: int
    categorical: bool = False

class DatasetEncoder:
    """Portable standard encoding recipe for an immutable dataset snapshot."""

    kind: Literal["rate_poisson", "prebinned_spikes", "event_bin"]
    duration_ms: float | None = None
    max_rate_hz: float | None = None
    seed: int = 0

class DatasetSnapshotBinding:
    """Bind one digest-identified NPZ snapshot to a graph input and target."""

    path: Path
    input_id: str
    dataset_id: str
    split: str
    encoder: DatasetEncoder
    target_id: str | None = None
    feature_key: str = "features"
    label_key: str = "labels"
    sample_cap: int | None = None
    shuffle: bool = False
    order_seed: int = 0

class ResolvedDenseInputs:
    tensors: Mapping[str, torch.Tensor]
    protocol: Mapping[str, Any]

class ExecutionSpec:
    kind: RequestKind
    executor: ExecutorName = "legacy"
    bundle: Path | None = None
    graph: Mapping[str, Any] | None = None
    inputs: Mapping[str, torch.Tensor] = field(default_factory=dict)
    input_bindings: Sequence[DenseArrayBinding] = field(default_factory=tuple)
    event_bindings: Sequence[EventStreamBinding] = field(default_factory=tuple)
    poisson_bindings: Sequence[PoissonInputBinding] = field(default_factory=tuple)
    dataset_binding: DatasetSnapshotBinding | None = None
    protocol: Mapping[str, Any] = field(default_factory=dict)
    training: Mapping[str, Any] | None = None
    targets: Mapping[str, torch.Tensor] = field(default_factory=dict)
    target_bindings: Sequence[TargetArrayBinding] = field(default_factory=tuple)
    seed: int = 0
    device: str = "auto"
    recording: RecordingProfile = "full"
    recording_fields: Sequence[str] | None = None
    checkpoint: Path | None = None
    runtime_state: GraphRuntimeState | None = None
    options: Mapping[str, Any] = field(default_factory=dict)

class ExecutionResult:
    executor: ExecutorName
    outputs: dict[str, torch.Tensor] = field(default_factory=dict)
    recordings: dict[str, torch.Tensor] = field(default_factory=dict)
    parameters: dict[str, torch.Tensor] = field(default_factory=dict)
    gradients: dict[str, torch.Tensor] = field(default_factory=dict)
    optimizer_state: dict[str, Any] = field(default_factory=dict)
    training_checkpoint: TrainingCheckpoint | None = None
    selected_checkpoint: TrainingCheckpoint | None = None
    final_state: dict[str, torch.Tensor] = field(default_factory=dict)
    runtime_state: GraphRuntimeState | None = None
    metrics: dict[str, Any] = field(default_factory=dict)
    model: nn.Module | None = None

class CapabilityIssue:
    element: str
    capability: str
    message: str

class TrainingCheckpoint:
    graph_digest: str
    training_digest: str
    completed_updates: int
    execution_protocol: Mapping[str, Any]
    initialization: Mapping[str, Any]
    parameters: dict[str, torch.Tensor]
    optimizer_state: dict[str, dict[str, Any]]
    rng_state: torch.Tensor
    rng_backend: str = "cpu"
    accelerator_rng_states: dict[str, torch.Tensor] = field(default_factory=dict)
    data_state: Mapping[str, Any] = field(default_factory=dict)
    selected_loss: float | None = None

class ParameterInterchange:
    parameters: dict[str, torch.Tensor]
    provenance: Mapping[str, Any]

def load_dense_array_bindings(
    path: str | Path, graph: Mapping[str, Any]
) -> tuple[DenseArrayBinding, ...]:
    """Load a replayable NPY/NPZ file without weakening named-input semantics."""
    source_path = Path(path)
    digest = "sha256:" + hashlib.sha256(source_path.read_bytes()).hexdigest()
    loaded = np.load(source_path, allow_pickle=False)
    input_ids = [row["id"] for row in graph.get("inputs", [])]
    source_base = {
        "kind": "file",
        "path": str(source_path),
        "digest": digest,
    }
    if isinstance(loaded, np.ndarray):
        if len(input_ids) != 1:
            raise ValueError(
                "a dense NPY can bind only a graph with exactly one input; "
                f"graph inputs are {input_ids}"
            )
        return (
            DenseArrayBinding(
                input_ids[0],
                torch.as_tensor(loaded),
                {**source_base, "array": None},
            ),
        )
    try:
        arrays = {key: loaded[key] for key in loaded.files}
    finally:
        loaded.close()
    if set(arrays) == {"input_spikes"} and len(input_ids) == 1:
        arrays = {input_ids[0]: arrays["input_spikes"]}
        keys = {input_ids[0]: "input_spikes"}
    else:
        keys = {key: key for key in arrays}
    return tuple(
        DenseArrayBinding(
            input_id,
            torch.as_tensor(value),
            {**source_base, "array": keys[input_id]},
        )
        for input_id, value in arrays.items()
    )

def load_target_array_bindings(
    path: str | Path, training: Mapping[str, Any]
) -> tuple[TargetArrayBinding, ...]:
    """Load NPY/NPZ class targets against the recipe's named objectives."""
    source_path = Path(path)
    digest = "sha256:" + hashlib.sha256(source_path.read_bytes()).hexdigest()
    target_ids = sorted(
        {row["target"] for row in training.get("objectives", []) if "target" in row}
    )
    loaded = np.load(source_path, allow_pickle=False)
    source = {"kind": "file", "path": str(source_path), "digest": digest}
    if isinstance(loaded, np.ndarray):
        if len(target_ids) != 1:
            raise ValueError(
                f"a target NPY requires exactly one recipe target; targets are {target_ids}"
            )
        arrays = {target_ids[0]: loaded}
        keys = {target_ids[0]: None}
    else:
        try:
            arrays = {key: loaded[key] for key in loaded.files}
        finally:
            loaded.close()
        keys = {key: key for key in arrays}
    if set(arrays) != set(target_ids):
        raise ValueError(
            f"target array ids do not match recipe targets; expected={target_ids}, got={sorted(arrays)}"
        )
    return tuple(
        TargetArrayBinding(
            target_id,
            torch.as_tensor(arrays[target_id]),
            {**source, "array": keys[target_id]},
        )
        for target_id in target_ids
    )

def resolve_target_array_bindings(
    training: Mapping[str, Any],
    *,
    bindings: Sequence[TargetArrayBinding] = (),
    targets: Mapping[str, torch.Tensor] | None = None,
    sample_count: int,
    device: str | torch.device = "cpu",
) -> tuple[dict[str, torch.Tensor], list[dict[str, Any]]]:
    """Validate named one-dimensional integer targets and retain their provenance."""
    if bindings and targets:
        raise ValueError("provide target bindings or target tensors, not both")
    if not bindings:
        bindings = tuple(
            TargetArrayBinding(name, value, {"kind": "memory"})
            for name, value in (targets or {}).items()
        )
    expected = {
        row["target"] for row in training.get("objectives", []) if "target" in row
    }
    by_name = {binding.target_id: binding for binding in bindings}
    if len(by_name) != len(bindings):
        raise ValueError("duplicate target array binding")
    if set(by_name) != expected:
        raise ValueError(
            f"target ids do not match recipe; missing={sorted(expected - set(by_name))}, unexpected={sorted(set(by_name) - expected)}"
        )
    resolved = {}
    rows = []
    integer_dtypes = {torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8}
    for target_id in sorted(by_name):
        binding = by_name[target_id]
        value = binding.value
        if value.ndim != 1 or value.shape[0] != sample_count:
            raise ValueError(
                f"training target {target_id} expected shape [{sample_count}], got {list(value.shape)}"
            )
        if value.dtype not in integer_dtypes:
            raise ValueError(f"training target {target_id} must use an integer dtype")
        resolved[target_id] = value.to(device=device, dtype=torch.long)
        rows.append(
            {
                "target_id": target_id,
                "shape": list(value.shape),
                "dtype": str(value.dtype).removeprefix("torch."),
                "source": dict(binding.source),
            }
        )
    return resolved, rows

def load_event_stream_bindings(
    path: str | Path, graph: Mapping[str, Any]
) -> tuple[EventStreamBinding, ...]:
    """Load named sparse spike coordinates from a replayable NPZ file."""
    source_path = Path(path)
    if source_path.suffix.lower() != ".npz":
        raise ValueError("event-stream replay requires an NPZ file")
    digest = "sha256:" + hashlib.sha256(source_path.read_bytes()).hexdigest()
    loaded = np.load(source_path, allow_pickle=False)
    try:
        arrays = {key: loaded[key] for key in loaded.files}
    finally:
        loaded.close()
    input_ids = [row["id"] for row in graph.get("inputs", [])]
    fields = ("steps", "batches", "channels", "steps_count", "batch_size")
    plain = set(fields)
    use_plain = len(input_ids) == 1 and set(arrays) == plain
    expected = (
        plain
        if use_plain
        else {f"{input_id}.{field}" for input_id in input_ids for field in fields}
    )
    if set(arrays) != expected:
        raise ValueError(
            "event-stream NPZ keys do not match graph inputs; "
            f"expected={sorted(expected)}, got={sorted(arrays)}"
        )
    source_base = {"kind": "file", "path": str(source_path), "digest": digest}
    bindings = []
    for input_id in input_ids:

        def key(field: str) -> str:
            return field if use_plain else f"{input_id}.{field}"

        steps_count_value = np.asarray(arrays[key("steps_count")])
        batch_size_value = np.asarray(arrays[key("batch_size")])
        if steps_count_value.size != 1 or batch_size_value.size != 1:
            raise ValueError(
                f"event input {input_id} steps_count and batch_size must be scalars"
            )
        if not np.issubdtype(steps_count_value.dtype, np.integer) or not np.issubdtype(
            batch_size_value.dtype, np.integer
        ):
            raise ValueError(
                f"event input {input_id} steps_count and batch_size must use integer dtypes"
            )
        bindings.append(
            EventStreamBinding(
                input_id=input_id,
                steps=torch.as_tensor(arrays[key("steps")]),
                batches=torch.as_tensor(arrays[key("batches")]),
                channels=torch.as_tensor(arrays[key("channels")]),
                steps_count=int(steps_count_value.item()),
                batch_size=int(batch_size_value.item()),
                source={
                    **source_base,
                    "arrays": {field: key(field) for field in fields},
                },
            )
        )
    return tuple(bindings)

def resolve_dense_array_bindings(
    graph: Mapping[str, Any],
    *,
    bindings: Sequence[DenseArrayBinding] = (),
    inputs: Mapping[str, torch.Tensor] | None = None,
    device: str | torch.device = "cpu",
    seed: int = 0,
    protocol: Mapping[str, Any] | None = None,
) -> ResolvedDenseInputs:
    """Validate dense arrays, resolve symbolic axes, and freeze run provenance."""
    if bindings and inputs:
        raise ValueError("provide dense input bindings or input tensors, not both")
    if not bindings:
        bindings = tuple(
            DenseArrayBinding(name, value, {"kind": "memory"})
            for name, value in (inputs or {}).items()
        )
    specs = {row["id"]: row for row in graph.get("inputs", [])}
    by_name: dict[str, DenseArrayBinding] = {}
    for binding in bindings:
        if binding.input_id in by_name:
            raise ValueError(f"duplicate dense binding for input {binding.input_id}")
        by_name[binding.input_id] = binding
    if set(by_name) != set(specs):
        missing = sorted(set(specs) - set(by_name))
        unexpected = sorted(set(by_name) - set(specs))
        raise ValueError(
            f"dense input ids do not match graph inputs; missing={missing}, unexpected={unexpected}"
        )

    resolved: dict[str, torch.Tensor] = {}
    rows: list[dict[str, Any]] = []
    leading_shape: tuple[int, int] | None = None
    masks: list[str] = []
    for input_id in sorted(specs):
        spec = specs[input_id]
        binding = by_name[input_id]
        value = binding.value
        declared = spec.get("shape", [])
        if len(declared) < 2 or declared[:2] != ["time", "batch"]:
            raise ValueError(
                f"input {input_id} dense binding requires declared shape beginning with ['time', 'batch']"
            )
        if value.ndim == len(declared) - 1 and declared[1] == "batch":
            value = value.unsqueeze(1)
        if value.ndim != len(declared):
            raise ValueError(
                f"input {input_id} rank expected {len(declared)}, got {value.ndim}"
            )
        expected_tail = tuple(int(axis) for axis in declared[2:])
        if tuple(value.shape[2:]) != expected_tail:
            raise ValueError(
                f"input {input_id} trailing shape expected {expected_tail}, got {tuple(value.shape[2:])}"
            )
        current_leading = (int(value.shape[0]), int(value.shape[1]))
        if leading_shape is None:
            leading_shape = current_leading
        elif current_leading != leading_shape:
            raise ValueError(
                f"input {input_id} leading shape expected {leading_shape}, got {current_leading}"
            )
        signal_type = spec.get("signal_type")
        if signal_type == "mask":
            if value.dtype != torch.bool:
                if not torch.all((value == 0) | (value == 1)):
                    raise ValueError(
                        f"input {input_id} mask values must be boolean or zero/one"
                    )
                value = value.bool()
            masks.append(input_id)
        elif not (value.is_floating_point() or value.dtype == torch.bool):
            value = value.float()
        if value.is_floating_point() and not torch.isfinite(value).all():
            raise ValueError(f"input {input_id} contains non-finite values")
        if signal_type == "spikes" and not torch.all((value == 0) | (value == 1)):
            raise ValueError(
                f"input {input_id} spike values must be boolean or zero/one"
            )
        if signal_type != "mask":
            value = value.float()
        value = value.to(device)
        resolved[input_id] = value
        rows.append(
            {
                "input_id": input_id,
                "representation": "dense_array",
                "shape": list(value.shape),
                "dtype": str(value.dtype).removeprefix("torch."),
                "signal_type": signal_type,
                "unit": spec.get("unit"),
                "source": dict(binding.source),
            }
        )
    assert leading_shape is not None
    dt_ms = float(graph["timebase"]["dt"]["value"])
    supplied = dict(protocol or {})
    dataset = dict(supplied.pop("dataset", {}))
    reserved = {
        "schema",
        "binding_schema",
        "representation",
        "inputs",
        "timing",
        "masks",
        "seeds",
    }
    if reserved & supplied.keys():
        raise ValueError(
            f"execution protocol cannot override reserved fields {sorted(reserved & supplied.keys())}"
        )
    dataset.setdefault("identity", None)
    dataset.setdefault("split", None)
    dataset.setdefault("sample_cap", leading_shape[1])
    dataset.setdefault("batch_size", leading_shape[1])
    dataset.setdefault("shuffle", None)
    execution_protocol = {
        "schema": EXECUTION_PROTOCOL_SCHEMA,
        "binding_schema": DENSE_ARRAY_BINDING_SCHEMA,
        "representation": "dense_array",
        "inputs": rows,
        "dataset": dataset,
        "timing": {
            "dt_ms": dt_ms,
            "steps": leading_shape[0],
            "duration_ms": leading_shape[0] * dt_ms,
        },
        "masks": masks,
        "seeds": {"execution": int(seed)},
        **supplied,
    }
    try:
        json.dumps(execution_protocol, sort_keys=True)
    except TypeError as exc:
        raise ValueError(
            f"execution protocol must be JSON-serializable: {exc}"
        ) from exc
    return ResolvedDenseInputs(tensors=resolved, protocol=execution_protocol)

def resolve_event_stream_bindings(
    graph: Mapping[str, Any],
    *,
    bindings: Sequence[EventStreamBinding],
    device: str | torch.device = "cpu",
    seed: int = 0,
    protocol: Mapping[str, Any] | None = None,
) -> ResolvedDenseInputs:
    """Validate sparse spike coordinates and materialize binary graph inputs."""
    specs = {row["id"]: row for row in graph.get("inputs", [])}
    by_name: dict[str, EventStreamBinding] = {}
    for binding in bindings:
        if binding.input_id in by_name:
            raise ValueError(f"duplicate event binding for input {binding.input_id}")
        by_name[binding.input_id] = binding
    if set(by_name) != set(specs):
        missing = sorted(set(specs) - set(by_name))
        unexpected = sorted(set(by_name) - set(specs))
        raise ValueError(
            f"event input ids do not match graph inputs; missing={missing}, unexpected={unexpected}"
        )

    resolved: dict[str, torch.Tensor] = {}
    rows: list[dict[str, Any]] = []
    leading_shape: tuple[int, int] | None = None
    integer_dtypes = {
        torch.int8,
        torch.int16,
        torch.int32,
        torch.int64,
        torch.uint8,
    }
    for input_id in sorted(specs):
        spec = specs[input_id]
        binding = by_name[input_id]
        declared = spec.get("shape", [])
        if declared[:2] != ["time", "batch"] or len(declared) != 3:
            raise ValueError(
                f"input {input_id} event binding requires shape ['time', 'batch', channels]"
            )
        if spec.get("signal_type") != "spikes":
            raise ValueError(
                f"input {input_id} event binding requires signal_type spikes"
            )
        channels_count = int(declared[2])
        if (
            not isinstance(binding.steps_count, Integral)
            or isinstance(binding.steps_count, bool)
            or not isinstance(binding.batch_size, Integral)
            or isinstance(binding.batch_size, bool)
        ):
            raise ValueError(
                f"event input {input_id} steps_count and batch_size must be integers"
            )
        steps_count = int(binding.steps_count)
        batch_size = int(binding.batch_size)
        if steps_count <= 0 or batch_size <= 0:
            raise ValueError(
                f"event input {input_id} steps_count and batch_size must be positive"
            )
        coordinates = (binding.steps, binding.batches, binding.channels)
        if any(value.ndim != 1 for value in coordinates):
            raise ValueError(
                f"event input {input_id} coordinates must be one-dimensional"
            )
        lengths = {int(value.numel()) for value in coordinates}
        if len(lengths) != 1:
            raise ValueError(f"event input {input_id} coordinate lengths must match")
        if any(value.dtype not in integer_dtypes for value in coordinates):
            raise ValueError(
                f"event input {input_id} coordinates must use integer dtypes"
            )
        steps = binding.steps.to(dtype=torch.int64, device="cpu")
        batches = binding.batches.to(dtype=torch.int64, device="cpu")
        channels = binding.channels.to(dtype=torch.int64, device="cpu")
        bounds = (
            ("step", steps, steps_count),
            ("batch", batches, batch_size),
            ("channel", channels, channels_count),
        )
        for label, values, upper in bounds:
            if torch.any(values < 0) or torch.any(values >= upper):
                raise ValueError(
                    f"event input {input_id} {label} coordinates must be in [0, {upper})"
                )
        flat = (steps * batch_size + batches) * channels_count + channels
        if flat.numel() > 1:
            differences = flat[1:] - flat[:-1]
            if torch.any(differences < 0):
                raise ValueError(
                    f"event input {input_id} coordinates must be ordered by step, batch, channel"
                )
            if torch.any(differences == 0):
                raise ValueError(
                    f"event input {input_id} contains duplicate coordinates"
                )
        current_leading = (steps_count, batch_size)
        if leading_shape is None:
            leading_shape = current_leading
        elif current_leading != leading_shape:
            raise ValueError(
                f"event input {input_id} leading shape expected {leading_shape}, got {current_leading}"
            )
        value = torch.zeros(
            (steps_count, batch_size, channels_count),
            dtype=torch.float32,
            device=device,
        )
        if flat.numel():
            value[
                steps.to(device=device),
                batches.to(device=device),
                channels.to(device=device),
            ] = 1.0
        resolved[input_id] = value
        rows.append(
            {
                "input_id": input_id,
                "representation": "event_stream",
                "shape": list(value.shape),
                "dtype": "float32",
                "signal_type": "spikes",
                "unit": spec.get("unit"),
                "event_count": int(flat.numel()),
                "source": dict(binding.source),
            }
        )
    if leading_shape is None:
        raise ValueError("graph execution requires at least one event input binding")
    dt_ms = float(graph["timebase"]["dt"]["value"])
    supplied = dict(protocol or {})
    dataset = dict(supplied.pop("dataset", {}))
    reserved = {
        "schema",
        "binding_schema",
        "representation",
        "inputs",
        "dataset",
        "timing",
        "masks",
        "seeds",
        "resolution",
    }
    if reserved & supplied.keys():
        raise ValueError(
            f"execution protocol cannot override reserved fields {sorted(reserved & supplied.keys())}"
        )
    dataset.setdefault("identity", None)
    dataset.setdefault("split", None)
    dataset.setdefault("sample_cap", leading_shape[1])
    dataset.setdefault("batch_size", leading_shape[1])
    dataset.setdefault("shuffle", None)
    execution_protocol = {
        "schema": EXECUTION_PROTOCOL_SCHEMA,
        "binding_schema": EVENT_STREAM_BINDING_SCHEMA,
        "representation": "event_stream",
        "inputs": rows,
        "dataset": dataset,
        "timing": {
            "dt_ms": dt_ms,
            "steps": leading_shape[0],
            "duration_ms": leading_shape[0] * dt_ms,
        },
        "masks": [],
        "seeds": {"execution": int(seed)},
        "resolution": {
            "coordinates": "zero_based_integer_steps",
            "ordering": "step,batch,channel",
            "duplicates": "reject",
            "materialization": "binary_dense",
        },
        **supplied,
    }
    try:
        json.dumps(execution_protocol, sort_keys=True)
    except TypeError as exc:
        raise ValueError(
            f"execution protocol must be JSON-serializable: {exc}"
        ) from exc
    return ResolvedDenseInputs(tensors=resolved, protocol=execution_protocol)

def resolve_input_bindings(
    graph: Mapping[str, Any],
    *,
    dense_bindings: Sequence[DenseArrayBinding] = (),
    event_bindings: Sequence[EventStreamBinding] = (),
    poisson_bindings: Sequence[PoissonInputBinding] = (),
    inputs: Mapping[str, torch.Tensor] | None = None,
    device: str | torch.device = "cpu",
    seed: int = 0,
    protocol: Mapping[str, Any] | None = None,
) -> ResolvedDenseInputs:
    """Resolve dense, event-stream, generated-Poisson, or mixed graph inputs."""
    if dense_bindings and inputs:
        raise ValueError("provide dense input bindings or input tensors, not both")
    if inputs:
        dense_bindings = tuple(
            DenseArrayBinding(name, value, {"kind": "memory"})
            for name, value in inputs.items()
        )
    dense_ids = {binding.input_id for binding in dense_bindings}
    event_ids = {binding.input_id for binding in event_bindings}
    poisson_ids = {binding.input_id for binding in poisson_bindings}
    overlap = sorted(
        (dense_ids & event_ids) | (dense_ids & poisson_ids) | (event_ids & poisson_ids)
    )
    if overlap:
        raise ValueError(
            f"graph inputs cannot have dense and event bindings: {overlap}"
        )
    graph_ids = {row["id"] for row in graph.get("inputs", [])}
    if dense_ids | event_ids | poisson_ids != graph_ids:
        missing = sorted(graph_ids - dense_ids - event_ids - poisson_ids)
        unexpected = sorted((dense_ids | event_ids | poisson_ids) - graph_ids)
        raise ValueError(
            f"input ids do not match graph inputs; missing={missing}, unexpected={unexpected}"
        )
    if poisson_bindings:
        if dense_bindings or event_bindings:
            raise ValueError(
                "Poisson bindings cannot yet be mixed with replay bindings"
            )
        return resolve_poisson_input_bindings(
            graph,
            bindings=poisson_bindings,
            device=device,
            seed=seed,
            protocol=protocol,
        )
    if not event_bindings:
        return resolve_dense_array_bindings(
            graph,
            bindings=dense_bindings,
            device=device,
            seed=seed,
            protocol=protocol,
        )
    if not dense_bindings:
        return resolve_event_stream_bindings(
            graph,
            bindings=event_bindings,
            device=device,
            seed=seed,
            protocol=protocol,
        )

    def graph_with_inputs(input_ids: set[str]) -> dict[str, Any]:
        return {
            **graph,
            "inputs": [
                row for row in graph.get("inputs", []) if row["id"] in input_ids
            ],
        }

    dense = resolve_dense_array_bindings(
        graph_with_inputs(dense_ids),
        bindings=dense_bindings,
        device=device,
        seed=seed,
        protocol=protocol,
    )
    events = resolve_event_stream_bindings(
        graph_with_inputs(event_ids),
        bindings=event_bindings,
        device=device,
        seed=seed,
        protocol=protocol,
    )
    if dense.protocol["timing"] != events.protocol["timing"]:
        raise ValueError(
            "dense and event input bindings must resolve to the same timestep, duration, and batch shape"
        )
    execution_protocol = {
        "schema": EXECUTION_PROTOCOL_SCHEMA,
        "binding_schema": MIXED_INPUT_BINDING_SCHEMA,
        "representation": "mixed",
        "inputs": sorted(
            [*dense.protocol["inputs"], *events.protocol["inputs"]],
            key=lambda row: row["input_id"],
        ),
        "dataset": dense.protocol["dataset"],
        "timing": dense.protocol["timing"],
        "masks": dense.protocol["masks"],
        "seeds": dense.protocol["seeds"],
        "resolution": {"event_stream": events.protocol["resolution"]},
        **{
            key: value
            for key, value in dense.protocol.items()
            if key
            not in {
                "schema",
                "binding_schema",
                "representation",
                "inputs",
                "dataset",
                "timing",
                "masks",
                "seeds",
                "resolution",
            }
        },
    }
    return ResolvedDenseInputs(
        tensors={**dense.tensors, **events.tensors}, protocol=execution_protocol
    )

def resolve_poisson_input_bindings(
    graph: Mapping[str, Any],
    *,
    bindings: Sequence[PoissonInputBinding],
    device: str | torch.device = "cpu",
    seed: int = 0,
    protocol: Mapping[str, Any] | None = None,
) -> ResolvedDenseInputs:
    """Generate reproducible fixed or per-presentation categorical Poisson spikes."""
    specs = {row["id"]: row for row in graph.get("inputs", [])}
    by_name = {binding.input_id: binding for binding in bindings}
    if len(by_name) != len(bindings):
        raise ValueError("duplicate Poisson binding for a graph input")
    if set(by_name) != set(specs):
        missing = sorted(set(specs) - set(by_name))
        unexpected = sorted(set(by_name) - set(specs))
        raise ValueError(
            f"Poisson input ids do not match graph inputs; missing={missing}, unexpected={unexpected}"
        )
    dt_ms = float(graph["timebase"]["dt"]["value"])
    tensors: dict[str, torch.Tensor] = {}
    rows: list[dict[str, Any]] = []
    leading_shape: tuple[int, int] | None = None
    for input_id in sorted(specs):
        spec = specs[input_id]
        binding = by_name[input_id]
        declared = spec.get("shape", [])
        if declared[:2] != ["time", "batch"] or len(declared) != 3:
            raise ValueError(
                f"input {input_id} Poisson binding requires shape ['time', 'batch', channels]"
            )
        if spec.get("signal_type") != "spikes":
            raise ValueError(
                f"input {input_id} Poisson binding requires signal_type spikes"
            )
        if binding.steps_count <= 0 or binding.batch_size <= 0:
            raise ValueError(
                f"Poisson input {input_id} steps_count and batch_size must be positive"
            )
        rates = tuple(float(rate) for rate in binding.rates_hz)
        if not rates or any(not math.isfinite(rate) or rate < 0 for rate in rates):
            raise ValueError(
                f"Poisson input {input_id} rates must be finite and non-negative"
            )
        if not binding.categorical and len(rates) != 1:
            raise ValueError(
                f"fixed-rate Poisson input {input_id} requires exactly one rate"
            )
        if max(rates) * dt_ms / 1000.0 > 1.0:
            raise ValueError(
                f"Poisson input {input_id} rate times dt exceeds probability one"
            )
        current_leading = (int(binding.steps_count), int(binding.batch_size))
        if leading_shape is None:
            leading_shape = current_leading
        elif current_leading != leading_shape:
            raise ValueError(
                f"Poisson input {input_id} leading shape expected {leading_shape}, got {current_leading}"
            )
        generator = torch.Generator(device="cpu").manual_seed(int(binding.seed))
        if binding.categorical:
            indices = torch.randint(
                len(rates), (binding.batch_size,), generator=generator
            )
            realized = torch.tensor(rates, dtype=torch.float32)[indices]
        else:
            realized = torch.full((binding.batch_size,), rates[0], dtype=torch.float32)
        probability = realized.reshape(1, -1, 1) * dt_ms / 1000.0
        value = (
            (
                torch.rand(
                    binding.steps_count,
                    binding.batch_size,
                    int(declared[2]),
                    generator=generator,
                )
                < probability
            )
            .float()
            .to(device)
        )
        tensors[input_id] = value
        rows.append(
            {
                "input_id": input_id,
                "representation": "poisson",
                "shape": list(value.shape),
                "dtype": "float32",
                "signal_type": "spikes",
                "unit": spec.get("unit"),
                "protocol": "categorical_rate" if binding.categorical else "fixed_rate",
                "rates_hz": list(rates),
                "realized_rates_hz": realized.tolist(),
                "seed": int(binding.seed),
                "selection": "uniform_independent_per_presentation"
                if binding.categorical
                else "constant",
            }
        )
    assert leading_shape is not None
    supplied = dict(protocol or {})
    dataset = dict(supplied.pop("dataset", {}))
    if {
        "schema",
        "binding_schema",
        "representation",
        "inputs",
        "timing",
        "seeds",
    } & supplied.keys():
        raise ValueError("execution protocol cannot override reserved Poisson fields")
    dataset.setdefault("identity", None)
    dataset.setdefault("split", None)
    dataset.setdefault("sample_cap", leading_shape[1])
    dataset.setdefault("batch_size", leading_shape[1])
    dataset.setdefault("shuffle", None)
    execution_protocol = {
        "schema": EXECUTION_PROTOCOL_SCHEMA,
        "binding_schema": POISSON_INPUT_BINDING_SCHEMA,
        "representation": "poisson",
        "inputs": rows,
        "dataset": dataset,
        "timing": {
            "dt_ms": dt_ms,
            "steps": leading_shape[0],
            "duration_ms": leading_shape[0] * dt_ms,
        },
        "masks": [],
        "seeds": {
            "execution": int(seed),
            "poisson": {row["input_id"]: row["seed"] for row in rows},
        },
        "resolution": {
            "distribution": "Bernoulli discretization of homogeneous Poisson",
            "rate_selection": "per_presentation",
        },
        **supplied,
    }
    json.dumps(execution_protocol, sort_keys=True)
    return ResolvedDenseInputs(tensors=tensors, protocol=execution_protocol)

def resolve_dataset_snapshot_binding(
    graph: Mapping[str, Any],
    binding: DatasetSnapshotBinding,
    *,
    device: str | torch.device = "cpu",
    execution_seed: int = 0,
    protocol: Mapping[str, Any] | None = None,
) -> tuple[ResolvedDenseInputs, tuple[TargetArrayBinding, ...]]:
    """Load, select, and encode one immutable MNIST/SHD-style NPZ snapshot."""
    specs = {row["id"]: row for row in graph.get("inputs", [])}
    if set(specs) != {binding.input_id}:
        raise ValueError(
            "dataset snapshot binding requires exactly its one named graph input"
        )
    spec = specs[binding.input_id]
    declared = spec.get("shape", [])
    if declared[:2] != ["time", "batch"] or len(declared) != 3:
        raise ValueError(
            "dataset snapshot binding requires graph input shape ['time', 'batch', channels]"
        )
    if spec.get("signal_type") != "spikes":
        raise ValueError("dataset snapshot binding requires a spike graph input")
    source_path = Path(binding.path)
    if source_path.suffix.lower() != ".npz":
        raise ValueError("dataset snapshot binding requires an NPZ file")
    source_digest = _file_digest(source_path)
    loaded = np.load(source_path, allow_pickle=False)
    try:
        arrays = {key: loaded[key] for key in loaded.files}
    finally:
        loaded.close()
    if not binding.dataset_id or not binding.split:
        raise ValueError("dataset snapshot identity and split must be non-empty")
    if binding.label_key not in arrays:
        raise ValueError(
            f"dataset snapshot is missing labels key {binding.label_key!r}"
        )
    labels = np.asarray(arrays[binding.label_key])
    if labels.ndim != 1 or not np.issubdtype(labels.dtype, np.integer):
        raise ValueError(
            "dataset snapshot labels must be a one-dimensional integer array"
        )
    sample_count = int(labels.shape[0])
    if sample_count <= 0:
        raise ValueError("dataset snapshot must contain at least one sample")
    cap = sample_count if binding.sample_cap is None else int(binding.sample_cap)
    if cap <= 0 or cap > sample_count:
        raise ValueError(f"dataset snapshot sample cap must be in [1, {sample_count}]")
    order = torch.arange(sample_count)
    if binding.shuffle:
        generator = torch.Generator(device="cpu").manual_seed(int(binding.order_seed))
        order = torch.randperm(sample_count, generator=generator)
    selected = order[:cap].numpy()
    selected_labels = labels[selected].astype(np.int64, copy=False)
    encoder = binding.encoder
    dt_ms = float(graph["timebase"]["dt"]["value"])
    channels = int(declared[2])
    encoder_row: dict[str, Any]
    if encoder.kind == "rate_poisson" and (
        encoder.duration_ms is None or encoder.max_rate_hz is None
    ):
        raise ValueError(
            "rate-Poisson encoder requires explicit duration_ms and max_rate_hz"
        )
    if encoder.kind == "prebinned_spikes" and (
        encoder.duration_ms is not None
        or encoder.max_rate_hz is not None
        or encoder.seed != 0
    ):
        raise ValueError(
            "prebinned encoder does not accept duration, max rate, or a stochastic seed"
        )
    if encoder.kind == "event_bin" and (
        encoder.duration_ms is None
        or encoder.max_rate_hz is not None
        or encoder.seed != 0
    ):
        raise ValueError(
            "event-bin encoder requires duration_ms and does not accept max rate or a stochastic seed"
        )
    if encoder.kind in {"rate_poisson", "prebinned_spikes"}:
        if binding.feature_key not in arrays:
            raise ValueError(
                f"dataset snapshot is missing feature key {binding.feature_key!r}"
            )
        features = np.asarray(arrays[binding.feature_key])
    if encoder.kind == "rate_poisson":
        if features.shape != (sample_count, channels):
            raise ValueError(
                f"rate-Poisson dataset features expected shape {(sample_count, channels)}, got {features.shape}"
            )
        if not np.issubdtype(features.dtype, np.floating):
            raise ValueError("rate-Poisson dataset features must use a floating dtype")
        if (
            not np.isfinite(features).all()
            or np.any(features < 0)
            or np.any(features > 1)
        ):
            raise ValueError("rate-Poisson dataset features must be finite in [0, 1]")
        duration_ms = float(encoder.duration_ms or 0)
        max_rate_hz = float(encoder.max_rate_hz or 0)
        raw_steps = duration_ms / dt_ms
        if (
            duration_ms <= 0
            or not math.isclose(raw_steps, round(raw_steps), abs_tol=1e-9)
            or not math.isfinite(max_rate_hz)
            or max_rate_hz < 0
            or max_rate_hz * dt_ms / 1000.0 > 1
        ):
            raise ValueError(
                "rate-Poisson encoder requires timestep-aligned positive duration and a supported finite non-negative max rate"
            )
        rates = torch.as_tensor(features[selected], dtype=torch.float32)
        generator = torch.Generator(device="cpu").manual_seed(int(encoder.seed))
        probability = rates.unsqueeze(0) * max_rate_hz * dt_ms / 1000.0
        spikes = (
            torch.rand(int(round(raw_steps)), cap, channels, generator=generator)
            < probability
        ).float()
        encoder_row = {
            "kind": encoder.kind,
            "duration_ms": duration_ms,
            "max_rate_hz": max_rate_hz,
            "seed": int(encoder.seed),
            "distribution": "Bernoulli discretization of feature-scaled homogeneous Poisson",
        }
    elif encoder.kind == "prebinned_spikes":
        if features.ndim != 3 or features.shape[1:] != (sample_count, channels):
            raise ValueError(
                f"prebinned dataset features expected [time, {sample_count}, {channels}]"
            )
        if not np.all((features == 0) | (features == 1)):
            raise ValueError("prebinned dataset spikes must be binary")
        spikes = torch.as_tensor(features[:, selected, :], dtype=torch.float32)
        encoder_row = {
            "kind": encoder.kind,
            "duration_ms": int(features.shape[0]) * dt_ms,
            "seed": None,
        }
    elif encoder.kind == "event_bin":
        event_keys = ("event_sample", "event_time_ms", "event_channel")
        missing = [key for key in event_keys if key not in arrays]
        if missing:
            raise ValueError(f"event dataset snapshot is missing keys {missing}")
        samples = np.asarray(arrays["event_sample"])
        times = np.asarray(arrays["event_time_ms"])
        event_channels = np.asarray(arrays["event_channel"])
        if (
            samples.ndim != 1
            or times.ndim != 1
            or event_channels.ndim != 1
            or not (len(samples) == len(times) == len(event_channels))
            or not np.issubdtype(samples.dtype, np.integer)
            or not np.issubdtype(event_channels.dtype, np.integer)
            or not np.issubdtype(times.dtype, np.floating)
        ):
            raise ValueError(
                "event dataset coordinates must be equal-length sample/channel integer and time floating arrays"
            )
        duration_ms = float(encoder.duration_ms or 0)
        raw_steps = duration_ms / dt_ms
        if duration_ms <= 0 or not math.isclose(
            raw_steps, round(raw_steps), abs_tol=1e-9
        ):
            raise ValueError(
                "event-bin encoder duration must be a positive integer number of timesteps"
            )
        if (
            not np.isfinite(times).all()
            or np.any(samples < 0)
            or np.any(samples >= sample_count)
            or np.any(event_channels < 0)
            or np.any(event_channels >= channels)
            or np.any(times < 0)
            or np.any(times >= duration_ms)
        ):
            raise ValueError("event dataset coordinates are out of snapshot bounds")
        selected_position = {
            int(sample): index for index, sample in enumerate(selected)
        }
        spikes = torch.zeros(int(round(raw_steps)), cap, channels)
        retained = 0
        collisions = 0
        for sample, time_ms, channel in zip(
            samples, times, event_channels, strict=True
        ):
            batch = selected_position.get(int(sample))
            if batch is None:
                continue
            step = min(int(float(time_ms) / dt_ms), spikes.shape[0] - 1)
            collisions += int(spikes[step, batch, int(channel)] != 0)
            spikes[step, batch, int(channel)] = 1
            retained += 1
        encoder_row = {
            "kind": encoder.kind,
            "duration_ms": duration_ms,
            "seed": None,
            "timestamp_unit": "ms",
            "binning": "floor_left_closed_right_open",
            "retained_events": retained,
            "binary_collisions": collisions,
        }
    else:
        raise ValueError(f"unsupported dataset encoder {encoder.kind!r}")
    source = {
        "kind": "dataset_snapshot",
        "path": str(source_path),
        "digest": source_digest,
        "arrays": {
            "labels": binding.label_key,
            **(
                {"features": binding.feature_key}
                if encoder.kind != "event_bin"
                else {
                    "sample": "event_sample",
                    "time_ms": "event_time_ms",
                    "channel": "event_channel",
                }
            ),
        },
    }
    resolved = resolve_dense_array_bindings(
        graph,
        bindings=(DenseArrayBinding(binding.input_id, spikes, source),),
        device=device,
        seed=execution_seed,
        protocol={
            "dataset": {
                "identity": binding.dataset_id,
                "split": binding.split,
                "sample_cap": cap,
                "batch_size": cap,
                "shuffle": bool(binding.shuffle),
            }
        },
    )
    supplied = dict(protocol or {})
    supplied_dataset = dict(supplied.pop("dataset", {}))
    expected_dataset = dict(resolved.protocol["dataset"])
    unknown_dataset = sorted(set(supplied_dataset) - set(expected_dataset))
    conflicts = sorted(
        key
        for key, value in supplied_dataset.items()
        if value is not None and value != expected_dataset[key]
    )
    if unknown_dataset or conflicts:
        raise ValueError(
            "dataset execution protocol metadata does not match binding; "
            f"unknown={unknown_dataset}, conflicts={conflicts}"
        )
    reserved = {
        "schema",
        "binding_schema",
        "representation",
        "inputs",
        "timing",
        "masks",
        "seeds",
        "dataset_binding",
    }
    if reserved & supplied.keys():
        raise ValueError(
            f"dataset execution protocol cannot override reserved fields {sorted(reserved & supplied.keys())}"
        )
    protocol = {
        **resolved.protocol,
        "binding_schema": DATASET_SNAPSHOT_BINDING_SCHEMA,
        "representation": "dataset_snapshot",
        "dataset_binding": {
            "source": source,
            "sample_count": sample_count,
            "selected_indices": selected.tolist(),
            "order_seed": int(binding.order_seed),
            "encoder": encoder_row,
            "target_id": binding.target_id,
        },
        "seeds": {
            **resolved.protocol["seeds"],
            "dataset_order": int(binding.order_seed),
            "encoder": int(encoder.seed) if encoder.kind == "rate_poisson" else None,
        },
        **supplied,
    }
    targets = (
        (
            TargetArrayBinding(
                binding.target_id,
                torch.as_tensor(selected_labels),
                source={**source, "array": binding.label_key},
            ),
        )
        if binding.target_id
        else ()
    )
    return ResolvedDenseInputs(resolved.tensors, protocol), targets