engine/io_utils.py

"""noctilith/io_utils.py — serialization helpers."""
from __future__ import annotations

import hashlib
import json
import pathlib
from typing import Any, Dict, Mapping

import numpy as np

SCHEMA_VERSION = "noctilith.schema.v1"


def _json_default(value: Any) -> Any:
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, (np.floating, np.integer, np.bool_)):
        return value.item()
    if isinstance(value, pathlib.Path):
        return str(value)
    return str(value)


def save_json(path: str, data: Dict[str, Any], *, indent: int = 2) -> str:
    p = pathlib.Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(data, indent=indent, default=_json_default), encoding="utf-8")
    return str(p)


def load_json(path: str) -> Dict[str, Any]:
    return json.loads(pathlib.Path(path).read_text(encoding="utf-8"))


def schema_header(module_name: str, module_version: str = "1.0.0") -> Dict[str, str]:
    return {
        "schema_version": SCHEMA_VERSION,
        "module_name": module_name,
        "module_version": module_version,
    }


def compute_state_hash(data: Mapping[str, Any]) -> str:
    """Stable hash for nested snapshot-like mappings."""
    payload = json.dumps(data, sort_keys=True, separators=(",", ":"), default=_json_default)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()



def save_snapshot_json(path: str, snapshot: Mapping[str, Any], *, indent: int = 2) -> str:
    return save_json(path, dict(snapshot), indent=indent)



def save_snapshot_npz(path: str, snapshot: Mapping[str, Any]) -> str:
    p = pathlib.Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)

    arrays: Dict[str, np.ndarray] = {}
    metadata: Dict[str, Any] = {}

    for key, value in snapshot.items():
        if isinstance(value, np.ndarray):
            arrays[key] = value
        else:
            metadata[key] = _json_default(value)

    arrays["__metadata_json__"] = np.array(json.dumps(metadata, sort_keys=True), dtype=object)
    np.savez_compressed(p, **arrays)
    return str(p)