palimseste/serialization.py

"""PALIMPSESTE — Serialization (save/load to disk).

The whole substrate is an append-only log of ``Trace`` records plus a small
amount of config/encoder state. Serialization is therefore straightforward:

  - Memory traces   -> a compact binary blob (one record per line: id, weight,
                       meta flag, packed bits of address + value)
  - Encoder atoms   -> JSON index + a packed-bits blob
  - Config          -> JSON

We avoid pickle for security/portability: everything is numpy arrays + JSON.
A saved model directory looks like::

    model_dir/
      config.json
      memory.bin          # traces (address + value packed bits, weights, flags)
      memory_index.bin    # LSH projection positions (for exact reload)
      encoder.json        # atom/role/level tables (keys -> ids)
      encoder.bin         # packed bits for all atom/role/level HVs
      tokenizer.json      # if a tokenizer is attached

This format is deterministic, auditable, and loads on any machine without
the original Python objects.
"""

from __future__ import annotations

from dataclasses import asdict
import json
import os
from pathlib import Path
import numpy as np

from .hv import HV, DEFAULT_D
from .memory import Memory, Trace
from .phi import Phi, KernelConfig
from .learner import Encoder
from .lsh import LSHConfig, LSHIndex

__all__ = ["save_memory", "load_memory", "save_encoder", "load_encoder",
           "save_config", "load_config"]


# ----------------------------------------------------------------- config
def save_config(path: str | Path, config: dict) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with open(p, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)


def load_config(path: str | Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


# ----------------------------------------------------------------- memory
def save_memory(mem: Memory, path: str | Path) -> None:
    """Serialize a :class:`Memory` to ``memory.bin`` + ``memory_index.bin``.

    Format of ``memory.bin`` (numpy ``.npy``-style, but custom for compactness):

      Header (JSON, one line):
        {"D": D, "n_traces": N, "n_meta": M, "decay": {...}, "version": 1}
      Body (raw bytes):
        For each trace (non-meta first, then meta):
          id (int64), weight (float64), t_insert (float64), tag_len (int32),
          tag (utf-8 bytes), address_bits (D/8 packed), value_bits (D/8 packed)
    """
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)

    header = {
        "D": mem.D,
        "n_traces": len(mem._traces),
        "n_meta": len(mem._meta_traces),
        "decay": mem.decay,
        "version": 1,
    }

    chunks: list[bytes] = []
    chunks.append((json.dumps(header) + "\n").encode("utf-8"))

    # Vectorized serialization: for large |M|, skip tags entirely (version 2
    # format) and use fixed-width records. Each record is:
    #   id(int64) + weight(float64) + t_insert(float64) + addr_bits + val_bits
    # = 24 + 2*packed_len bytes per trace, all stackable as a numpy array.
    packed_len = (mem.D + 7) // 8
    record_size = 24 + 2 * packed_len

    def _emit_batch_streaming(traces: list[Trace], f, chunk_size: int = 50000):
        """Write traces to file in chunks to avoid MemoryError on large |M|."""
        if not traces:
            return
        n = len(traces)
        for start in range(0, n, chunk_size):
            end = min(start + chunk_size, n)
            batch = traces[start:end]
            bn = len(batch)
            buf = np.empty((bn, record_size), dtype=np.uint8)
            ids = np.array([tr.id for tr in batch], dtype=np.int64)
            weights = np.array([tr.weight for tr in batch], dtype=np.float64)
            t_ins = np.array([tr.t_insert for tr in batch], dtype=np.float64)
            buf[:, :8] = np.frombuffer(ids.tobytes(), dtype=np.uint8).reshape(bn, 8)
            buf[:, 8:16] = np.frombuffer(weights.tobytes(), dtype=np.uint8).reshape(bn, 8)
            buf[:, 16:24] = np.frombuffer(t_ins.tobytes(), dtype=np.uint8).reshape(bn, 8)
            buf[:, 24:24+packed_len] = np.stack([tr.address.bits for tr in batch])
            buf[:, 24+packed_len:] = np.stack([tr.value.bits for tr in batch])
            f.write(buf.tobytes())

    with open(p, "wb") as f:
        f.write((json.dumps(header) + "\n").encode("utf-8"))
        _emit_batch_streaming(mem._traces, f)
        _emit_batch_streaming(mem._meta_traces, f)

    # Save LSH projection positions so the index is reproducible on reload.
    # np.save appends .npy, so we name the file accordingly.
    idx_path = p.with_name(p.stem + ".index.npy")
    assert mem._index is not None
    proj = mem._index._positions
    np.save(idx_path, proj)


def load_memory(path: str | Path, rng: np.random.Generator | None = None) -> Memory:
    """Load a :class:`Memory` from ``save_memory`` output.

    Rebuilds the LSH index with the *saved* projection positions so retrieval
    is bit-identical to the saved model.
    """
    p = Path(path)
    with open(p, "rb") as f:
        raw = f.read()

    # parse header
    nl = raw.index(b"\n")
    header = json.loads(raw[:nl].decode("utf-8"))
    D = header["D"]
    n_traces = header["n_traces"]
    n_meta = header["n_meta"]
    decay = header.get("decay", {"half_life": float("inf"), "floor": 1e-3})

    body = raw[nl + 1:]
    packed_len = (D + 7) // 8
    record_size = 24 + 2 * packed_len  # id(8) + weight(8) + t_insert(8) + addr + val

    # Vectorized read: interpret body as a (n_total, record_size) uint8 array
    n_total = n_traces + n_meta
    if n_total == 0:
        if rng is None:
            rng = np.random.default_rng()
        return Memory(D=D, decay=decay, rng=rng)

    buf = np.frombuffer(body, dtype=np.uint8, count=n_total * record_size).reshape(n_total, record_size)

    # Extract fixed fields
    ids = buf[:, :8].copy().view(np.int64).reshape(-1)
    weights = buf[:, 8:16].copy().view(np.float64).reshape(-1)
    t_ins = buf[:, 16:24].copy().view(np.float64).reshape(-1)
    addr_bits_all = buf[:, 24:24+packed_len].copy()
    val_bits_all = buf[:, 24+packed_len:].copy()

    if rng is None:
        rng = np.random.default_rng()

    mem = Memory(D=D, decay=decay, rng=rng)

    # Restore LSH projections BEFORE inserting traces
    idx_path = p.with_name(p.stem + ".index.npy")
    if idx_path.exists():
        proj = np.load(idx_path)
        assert mem._index is not None
        mem._index._positions = proj

    # Insert traces (non-meta first, then meta)
    for i in range(n_traces):
        tr = Trace(
            id=int(ids[i]),
            address=HV(bits=addr_bits_all[i], D=D),
            value=HV(bits=val_bits_all[i], D=D),
            weight=float(weights[i]),
            t_insert=float(t_ins[i]),
            meta=False,
            tag=None,
        )
        mem._traces.append(tr)
        assert mem._index is not None
        mem._index.insert(tr.id, tr.address)

    for i in range(n_traces, n_total):
        tr = Trace(
            id=int(ids[i]),
            address=HV(bits=addr_bits_all[i], D=D),
            value=HV(bits=val_bits_all[i], D=D),
            weight=float(weights[i]),
            t_insert=float(t_ins[i]),
            meta=True,
            tag=None,
        )
        mem._meta_traces.append(tr)
        mem._ensure_meta_index().insert(tr.id, tr.address)

    return mem


# ----------------------------------------------------------------- chunked I/O
def split_file(path: str | Path, chunk_size: int = 4_500_000_000,
               suffix: str = ".part") -> list[Path]:
    """Split a large binary file into chunks under ``chunk_size`` bytes.

    Creates files: ``path.part0``, ``path.part1``, ...
    Returns list of chunk paths.
    """
    p = Path(path)
    fsize = p.stat().st_size
    n_chunks = (fsize + chunk_size - 1) // chunk_size
    chunks: list[Path] = []

    with open(p, "rb") as f:
        for i in range(n_chunks):
            chunk_path = p.with_suffix(f"{suffix}{i}")
            remaining = min(chunk_size, fsize - i * chunk_size)
            with open(chunk_path, "wb") as out:
                written = 0
                while written < remaining:
                    block = f.read(min(1024 * 1024, remaining - written))
                    if not block:
                        break
                    out.write(block)
                    written += len(block)
            chunks.append(chunk_path)
            print(f"  {chunk_path.name}: {chunk_path.stat().st_size / 1024**3:.2f} GB")

    return chunks


def merge_files(chunk_paths: list[str | Path], output_path: str | Path) -> Path:
    """Merge chunk files back into a single file.

    Parameters
    ----------
    chunk_paths : list
        Ordered list of chunk file paths (part0, part1, ...).
    output_path : str | Path
        Where to write the merged file.
    """
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "wb") as f:
        for cp in chunk_paths:
            cp = Path(cp)
            with open(cp, "rb") as chunk:
                while True:
                    block = chunk.read(1024 * 1024)
                    if not block:
                        break
                    f.write(block)
    return out


def find_chunks(directory: str | Path, stem: str,
                suffix: str = ".part") -> list[Path]:
    """Find all chunk files for a given stem in a directory.

    Returns them sorted by part number.
    """
    d = Path(directory)
    chunks = sorted(d.glob(f"{stem}{suffix}*"))
    # Sort numerically by part index
    def _part_idx(p: Path) -> int:
        try:
            return int(p.name.rsplit(suffix, 1)[1])
        except (ValueError, IndexError):
            return 0
    return sorted(chunks, key=_part_idx)


def load_memory_chunked(directory: str | Path, base_name: str = "palimpseste_memory",
                        rng: np.random.Generator | None = None) -> Memory:
    """Load memory from chunked files.

    If ``directory/base_name.bin`` exists, loads directly.
    Otherwise, looks for ``directory/base_name.bin.part0``, ``.part1``, etc.,
    merges them to a temp file, then loads.
    """
    d = Path(directory)
    direct = d / f"{base_name}.bin"
    if direct.exists():
        return load_memory(direct, rng=rng)

    # Find chunks
    chunks = find_chunks(d, f"{base_name}.bin")
    if not chunks:
        # Try without .bin extension
        chunks = find_chunks(d, base_name)
    if not chunks:
        raise FileNotFoundError(f"No memory file or chunks found in {d} for {base_name}")

    print(f"Merging {len(chunks)} chunks...", flush=True)
    import tempfile
    with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp:
        tmp_path = Path(tmp.name)

    merge_files(chunks, tmp_path)
    mem = load_memory(tmp_path, rng=rng)
    tmp_path.unlink()  # cleanup temp
    return mem


# ----------------------------------------------------------------- encoder
def save_encoder(enc: Encoder, path: str | Path) -> None:
    """Serialize an :class:`Encoder` to JSON index + packed-bits blob."""
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)

    # Index: key -> row in the .bin blob
    atoms_keys = list(enc._atoms.keys())
    roles_keys = list(enc._roles.keys())

    header = {
        "D": enc.D,
        "n_atoms": len(atoms_keys),
        "n_roles": len(roles_keys),
        "n_levels": enc._n_levels,
        "has_levels": enc._levels is not None,
        "atoms": [_key_to_json(k) for k in atoms_keys],
        "roles": [_key_to_json(k) for k in roles_keys],
        "version": 1,
    }
    with open(p, "w", encoding="utf-8") as f:
        json.dump(header, f, indent=2, ensure_ascii=False)

    # Pack all HVs into one blob: atoms, then roles, then levels
    packed_len = (enc.D + 7) // 8
    rows: list[np.ndarray] = []
    for k in atoms_keys:
        rows.append(enc._atoms[k].bits)
    for k in roles_keys:
        rows.append(enc._roles[k].bits)
    if enc._levels is not None:
        for lv in enc._levels:
            rows.append(lv.bits)
    if rows:
        blob = np.stack(rows)  # (N, packed_len) uint8
    else:
        blob = np.zeros((0, packed_len), dtype=np.uint8)
    np.save(p.with_name(p.stem + ".bin.npy"), blob)


def load_encoder(path: str | Path, rng: np.random.Generator | None = None) -> Encoder:
    p = Path(path)
    with open(p, "r", encoding="utf-8") as f:
        header = json.load(f)
    D = header["D"]
    enc = Encoder(D=D, rng=rng if rng else np.random.default_rng())
    enc._n_levels = header["n_levels"]

    blob = np.load(p.with_name(p.stem + ".bin.npy"))
    packed_len = (D + 7) // 8
    row = 0
    for i in range(header["n_atoms"]):
        key = _key_from_json(header["atoms"][i])
        enc._atoms[key] = HV(bits=blob[row].copy(), D=D)
        row += 1
    for i in range(header["n_roles"]):
        key = _key_from_json(header["roles"][i])
        enc._roles[key] = HV(bits=blob[row].copy(), D=D)
        row += 1
    if header["has_levels"]:
        levels = []
        for _ in range(header["n_levels"]):
            levels.append(HV(bits=blob[row].copy(), D=D))
            row += 1
        enc._levels = levels
    return enc


# ----------------------------------------------------------------- helpers
def _key_to_json(k: object) -> list:
    """JSON-safe encoding of an encoder key.

    Encoder keys are either:
      - tuples like ('i', 5), ('s', 'cat'), ('b', True), ('__char__a')
      - plain ints (role positions: 0, 1, 2, ...)
    """
    if isinstance(k, tuple):
        # encode the elements: first element is a str tag, rest are scalars
        elems = [k[0]] + [e for e in k[1:]]
        return ["t"] + [_scalar_to_json(e) for e in elems]
    if isinstance(k, (int, np.integer)):
        return ["i", int(k)]
    return ["s", str(k)]


def _key_from_json(j: list) -> object:
    if j[0] == "t":
        elems = [_scalar_from_json(e) for e in j[1:]]
        return tuple(elems)
    if j[0] == "i":
        return int(j[1])
    return j[1]  # "s" -> str


def _scalar_to_json(e: object) -> object:
    if isinstance(e, (int, np.integer)):
        return {"i": int(e)}
    if isinstance(e, (float, np.floating)):
        return {"f": float(e)}
    if isinstance(e, bool):
        return {"b": e}
    return {"s": str(e)}


def _scalar_from_json(e: object) -> object:
    if isinstance(e, dict):
        if "i" in e:
            return int(e["i"])
        if "f" in e:
            return float(e["f"])
        if "b" in e:
            return bool(e["b"])
        if "s" in e:
            return str(e["s"])
    return e