# CelerisLab/common/checkpoint.py
"""HDF5 checkpoint save and load for CelerisLab simulations.

Responsibilities:
    Serialize GPU distribution functions, cell flags, simulation step count,
    and rigid-body states to HDF5; on load, validate lattice geometry and
    storage precision before restoring device buffers.

Usage::
    Call ``save_checkpoint`` / ``load_checkpoint`` from host-side code when GPU
    work is finished; the active ``LBMConfig`` must match the checkpoint on
    ``lattice_model``, ``nx``, ``ny``, ``nz``, and ``store_precision``.
"""

import json
import os
import time

import numpy as np
import pycuda.driver as cuda

# File format version — bump on incompatible layout changes.
CKPT_VERSION = 1
_SUPPORTED_STORE_PRECISIONS = ("FP32", "FP16S")


def _require_supported_precision(store_precision: str, source: str) -> None:
    supported = ", ".join(_SUPPORTED_STORE_PRECISIONS)
    if store_precision not in _SUPPORTED_STORE_PRECISIONS:
        raise RuntimeError(
            f"Unsupported store_precision '{store_precision}' in {source}. "
            f"Supported checkpoint/runtime precisions: {supported}."
        )


def save_checkpoint(field, stepper, lbm_cfg, bodies, path=None):
    """Save full simulation state to an HDF5 file.

    Args:
        field: LBMField instance (GPU arrays).
        stepper: LBMStepper instance (step count).
        lbm_cfg: LBMConfig instance.
        bodies: ObjectManager instance.
        path: Output file path. Auto-generated if None.

    Returns:
        Path to the written checkpoint file.
    """
    import h5py

    if path is None:
        os.makedirs("output", exist_ok=True)
        path = f"output/checkpoint_{stepper.step_count:08d}.h5"

    _require_supported_precision(lbm_cfg.store_precision, "save_checkpoint")
    # Download GPU → host
    dtype = np.float16 if field.store_bytes == 2 else np.float32
    ddf_buf = np.empty(field.n * field.nq, dtype=dtype)
    tmp_buf = np.empty(field.n * field.nq, dtype=dtype)
    cuda.memcpy_dtoh(ddf_buf, field.ddf_gpu)
    cuda.memcpy_dtoh(tmp_buf, field.temp_gpu)
    flag_buf = np.empty(field.n, dtype=np.uint16)
    cuda.memcpy_dtoh(flag_buf, field.flag_gpu)

    with h5py.File(path, "w") as hf:
        # Version and metadata
        hf.attrs["ckpt_version"] = CKPT_VERSION
        hf.attrs["celerislab_version"] = __import__(
            "CelerisLab").__version__
        hf.attrs["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%S")
        hf.attrs["step_count"] = stepper.step_count

        # Config as JSON string (human-readable in HDF5 viewers)
        cfg = lbm_cfg
        config_dict = {
            "lattice_model": cfg.lattice_model,
            "nx": cfg.nx, "ny": cfg.ny, "nz": cfg.nz,
            "data_type": cfg.data_type,
            "viscosity": cfg.viscosity, "velocity": cfg.velocity,
            "rho": cfg.rho,
            "collision": cfg.collision, "streaming": cfg.streaming,
            "store_precision": cfg.store_precision,
            "ddf_shifting": cfg.ddf_shifting,
            "les_enabled": cfg.les_enabled, "les_cs": cfg.les_cs,
            "les_closed_form": cfg.les_closed_form,
            "trt_magic_param": cfg.trt_magic_param,
            "inlet_profile": cfg.inlet_profile,
            "outlet_mode": cfg.outlet_mode,
            "omega_min": cfg.omega_min, "omega_max": cfg.omega_max,
        }
        hf.attrs["config_json"] = json.dumps(config_dict)

        # GPU arrays — gzip compression for portable, compact files
        hf.create_dataset("ddf", data=ddf_buf, compression="gzip",
                          compression_opts=4)
        hf.create_dataset("temp", data=tmp_buf, compression="gzip",
                          compression_opts=4)
        hf.create_dataset("flag", data=flag_buf, compression="gzip",
                          compression_opts=4)

        # Object states
        obj_states = []
        for obj in bodies.objects:
            s = obj.state
            obj_states.append({
                "obj_id": obj.obj_id,
                "type": obj.obj_type,
                "x": s.x, "y": s.y, "z": s.z, "theta": s.theta,
                "vx": s.vx, "vy": s.vy, "omega": s.omega,
                "radius": obj.radius,
            })
        hf.attrs["objects_json"] = json.dumps(obj_states)

    return path


def load_checkpoint(path, field, stepper, lbm_cfg, bodies):
    """Restore simulation state from an HDF5 checkpoint.

    The current config's lattice_model, nx, ny, nz, and
    store_precision must match the checkpoint.  Other parameters
    (collision, LES, …) may differ for parameter-study restarts.

    Args:
        path: HDF5 checkpoint file path.
        field: LBMField instance.
        stepper: LBMStepper instance.
        lbm_cfg: LBMConfig instance.
        bodies: ObjectManager instance.
    """
    import h5py

    _require_supported_precision(lbm_cfg.store_precision, "load_checkpoint")
    with h5py.File(path, "r") as hf:
        # Version check
        ver = hf.attrs.get("ckpt_version", 0)
        if ver != CKPT_VERSION:
            raise ValueError(
                f"Checkpoint version mismatch: file={ver}, "
                f"expected={CKPT_VERSION}. "
                f"Created by CelerisLab {hf.attrs.get('celerislab_version', '?')}")

        # Config compatibility check
        saved_cfg = json.loads(hf.attrs["config_json"])
        for key in ("lattice_model", "nx", "ny", "nz", "store_precision"):
            saved_val = saved_cfg.get(key)
            curr_val = getattr(lbm_cfg, key)
            if saved_val != curr_val:
                raise ValueError(
                    f"Config mismatch on '{key}': "
                    f"checkpoint={saved_val}, current={curr_val}")
        _require_supported_precision(
            str(saved_cfg.get("store_precision", "unknown")),
            "checkpoint file",
        )

        # Restore GPU arrays
        ddf_buf = hf["ddf"][:]
        tmp_buf = hf["temp"][:]
        flag_buf = hf["flag"][:]

        cuda.memcpy_htod(field.ddf_gpu, ddf_buf)
        cuda.memcpy_htod(field.temp_gpu, tmp_buf)
        cuda.memcpy_htod(field.flag_gpu, flag_buf)
        field.flag = flag_buf.copy()

        # Restore step count
        stepper._step_count = int(hf.attrs["step_count"])

        # Restore object states
        obj_data = json.loads(hf.attrs.get("objects_json", "[]"))
        for od in obj_data:
            oid = od["obj_id"]
            try:
                obj = bodies.get(oid)
                obj.state.x = od["x"]
                obj.state.y = od["y"]
                obj.state.z = float(od.get("z", 0.0))
                obj.state.theta = od["theta"]
                obj.state.vx = od["vx"]
                obj.state.vy = od["vy"]
                obj.state.omega = od["omega"]
            except KeyError:
                pass  # object not in current sim — skip
