"""Shared utilities for serialization and checkpointing."""

import numpy as np


def serialize_rng_state(state: dict) -> dict:
    """Convert numpy RNG state to JSON-serializable form."""
    result = {}
    for k, v in state.items():
        if isinstance(v, np.ndarray):
            result[k] = {"__ndarray__": True, "data": v.tolist(), "dtype": str(v.dtype)}
        elif isinstance(v, dict):
            result[k] = serialize_rng_state(v)
        else:
            result[k] = v
    return result


def deserialize_rng_state(state: dict) -> dict:
    """Reconstruct numpy RNG state from JSON."""
    result = {}
    for k, v in state.items():
        if isinstance(v, dict) and v.get("__ndarray__"):
            result[k] = np.array(v["data"], dtype=v["dtype"])
        elif isinstance(v, dict):
            result[k] = deserialize_rng_state(v)
        else:
            result[k] = v
    return result
