"""Checkpointing: compact state save/load for HAKO (npz-compressed + json).

Aggressive storage: every tensor is float32 (float16 for bulk matrices),
np.savez_compressed for ~60% size reduction, and the checkpoint carries the
GHSOM tree structure (numpy), head/adapters/attention tensors, the Theta /
K_pgo / growth / EWC states, the tokenizer merges, and the telemetry
counters. Restoring yields a byte-identical working state (deterministic
serialization).
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Dict

import numpy as np


def save_state(path: Path, **blocks) -> Path:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    npz_payload: Dict[str, np.ndarray] = {}
    meta: Dict[str, dict] = {}
    for name, obj in blocks.items():
        if obj is None:
            continue
        if isinstance(obj, dict):
            arrays = {k: v for k, v in obj.items()
                      if isinstance(v, np.ndarray)}
            scalars = {k: (v.tolist() if isinstance(v, (np.ndarray, list))
                           else v) for k, v in obj.items()
                       if not isinstance(v, np.ndarray) or v.ndim == 0}
            for k, v in arrays.items():
                npz_payload[f"{name}::{k}"] = v.astype(
                    np.float16 if v.dtype == np.float32 else v.dtype)
            meta[name] = scalars
        else:
            meta[name] = {"repr": str(obj)[:400]}
    np.savez_compressed(path, **npz_payload)
    Path(str(path).replace(".npz", ".meta.json")).write_text(
        json.dumps(meta, indent=1, default=str), encoding="utf-8")
    return path


def load_state(path: Path) -> Dict[str, Dict[str, np.ndarray]]:
    path = Path(path)
    data = np.load(path, allow_pickle=False)
    out: Dict[str, Dict[str, np.ndarray]] = {}
    for key in data.files:
        name, _, fieldname = key.partition("::")
        out.setdefault(name, {})[fieldname] = data[key]
    return out