dima.py

# src/dima/dima.py
from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from typing import Any, Dict, Optional, Tuple, Union

import numpy as np

import jax
import jax.numpy as jnp
from jax import random

from flax import serialization as flax_ser

from ann import ANNBackend, make_ann
from ddpm import DDPM
from dmap import DMAP
from gplm import GPLM


# ----------------------------
# Optional: Hugging Face Hub
# ----------------------------
try:
    from huggingface_hub import HfApi, HfFolder, upload_file, hf_hub_download  # type: ignore
    _HAS_HF = True
except Exception:
    _HAS_HF = False


_UNSET = object()


def _select_device(prefer: str = "auto"):
    """
    Safe JAX device selection.
    prefer: "auto" | "gpu" | "cpu"
    """
    prefer = (prefer or "auto").lower()
    devs = jax.devices()
    gpu = [d for d in devs if d.platform == "gpu"]
    cpu = [d for d in devs if d.platform == "cpu"]

    if prefer in ("auto", "gpu"):
        return gpu[0] if gpu else (cpu[0] if cpu else devs[0])
    if prefer == "cpu":
        return cpu[0] if cpu else devs[0]
    return gpu[0] if gpu else (cpu[0] if cpu else devs[0])


def _np_dtype_str(x) -> str:
    try:
        return str(np.dtype(x))
    except Exception:
        return "float32"


# ----------------------------
# Frozen inference-only models
# ----------------------------
class FrozenDMAP:
    """
    Inference-only Nyström DMAP embedder built from saved DMAP state.
    Uses kNN in ambient space against reference points.
    """

    def __init__(
        self,
        state: Dict[str, Any],
        *,
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
    ):
        self.k = int(state["k"])
        self.beta = float(state["beta"])
        self.β = self.beta
        self.alpha = float(state["alpha"])
        self.α = self.alpha
        self.eps = float(state["eps"])
        self.ε = self.eps
        self.dtype = np.dtype(state.get("dtype", "float32"))

        self.R_iX = np.ascontiguousarray(np.asarray(state["R_iX"]).astype(self.dtype, copy=False))
        self.qalpha_i = np.asarray(state["qalpha_i"], dtype=np.float64)
        self.qα_i = self.qalpha_i  # alias
        self.R_over_lam_ix = np.asarray(state["R_over_lam_ix"], dtype=np.float64)  # (Nref,d)
        self.R_over_λ_ix = self.R_over_lam_ix  # alias

        self.ann, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
        self.ann.build(self.R_iX)

    def __call__(self, R_aX: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
        R_aX = np.asarray(R_aX)
        single = (R_aX.ndim == 1)
        if single:
            R_aX = R_aX[None, :]

        R_aX = np.ascontiguousarray(R_aX.astype(self.dtype, copy=False))

        if batch_size is None:
            Z = self._embed(R_aX)
        else:
            out = []
            bs = int(batch_size)
            for s in range(0, R_aX.shape[0], bs):
                out.append(self._embed(R_aX[s : s + bs]))
            Z = np.vstack(out)

        return Z[0] if single else Z

    def _embed(self, R_aX: np.ndarray) -> np.ndarray:
        j_aK, D2_aK = self.ann.search(R_aX, self.k)  # (a,k)
        K_ai = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)

        q_a = np.maximum(K_ai.sum(axis=1), 1e-30)
        qalpha_a = np.maximum(np.power(q_a, self.alpha), 1e-30)

        qalpha_i = np.maximum(self.qalpha_i[j_aK], 1e-30)
        Kalpha_ai = K_ai / (qalpha_a[:, None] * qalpha_i)

        d_a = np.maximum(Kalpha_ai.sum(axis=1), 1e-30)
        P_ai = Kalpha_ai / d_a[:, None]

        R_over = self.R_over_lam_ix[j_aK, :]              # (a,k,d)
        Z_ax = (P_ai[:, :, None] * R_over).sum(axis=1)    # (a,d)
        return Z_ax


def _restore_gplm_as_object(
    state: Dict[str, Any],
    *,
    ann_backend: ANNBackend = "auto",
    ann_params: Optional[Dict[str, Any]] = None,
    n_jobs: int = -1,
) -> GPLM:
    """
    Rehydrate a GPLM instance from saved state WITHOUT retraining.
    This is deliberately done as a true GPLM instance so you also get GPLM.flow(...)
    (assuming your GPLM class implements .flow()).
    """
    obj = GPLM.__new__(GPLM)  # type: ignore

    obj.beta = float(state["beta"])
    obj.β = obj.beta
    obj.eps = float(state["eps"])
    obj.ε = obj.eps
    obj.pred_k = None if state.get("pred_k", None) is None else int(state["pred_k"])
    obj.pred_κ = obj.pred_k
    obj.dtype = np.dtype(state.get("dtype", "float32"))

    obj.mean_X = np.asarray(state["mean_X"], dtype=np.float64)
    obj.M_mX = np.asarray(state["M_mX"], dtype=np.float64)

    obj.lat_mean_x = np.asarray(state["lat_mean_x"], dtype=np.float64)
    obj.lat_std_x = np.asarray(state["lat_std_x"], dtype=np.float64)

    obj.Z_mx_w = np.ascontiguousarray(np.asarray(state["Z_mx_w"]).astype(np.float64, copy=False))
    obj.m = int(obj.Z_mx_w.shape[0])

    obj.ann_Z, _ = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
    obj.ann_Z.build(obj.Z_mx_w.astype(obj.dtype, copy=False))

    return obj


# ----------------------------
# Config
# ----------------------------
@dataclass
class DIMAConfig:
    d: int = 32
    beta: float = 1.0
    ddpm_device: str = "auto"   # "auto" | "cpu" | "gpu"
    version: str = "0.2.0"


# ----------------------------
# Main wrapper
# ----------------------------
class DIMA:
    """
    DIMA: DMAP encoder + (latent DDPM) + GPLM decoder.
    Public “user-facing” convention in this wrapper:
      - Raw DMAP coordinates are the *public latent* (np.ndarray): R_ax (a,d)
      - Normalized latents are the DDPM coordinates (jax/np):      Z_ax (a,d)
    Minimal user API (what you asked for):
        dima = DIMA(R_iX, d=20, beta=2.0)
        R_ax = dima(R_aX)        # encode ambient -> raw latents
        Q_aX = dima(R_ax)        # decode raw latents -> ambient
    You can still pass full dict overrides for any submodule:
        dima = DIMA(..., dmap_kwargs={...}, gplm_kwargs={...}, ddpm_kwargs={...})
    and you can also tweak the “headline” DDPM knobs directly in __init__ (below).
    """

    def __init__(
        self,
        R_iX: np.ndarray,
        *,
        # main knobs
        d: int = 32,
        beta: float = 1.0,
        # allow per-module override (if None -> uses global beta)
        dmap_beta: Optional[float] = None,
        gplm_beta: Optional[float] = None,
        # DMAP headline knobs (everything else via dmap_kwargs)
        dmap_alpha: float = 0.0,
        dmap_t: float = 1.0,
        dmap_k: Optional[int] = None,
        # GPLM headline knobs (everything else via gplm_kwargs)
        gplm_m: int = 1024,
        gplm_pred_k: Optional[int] = None,
        # DDPM headline knobs (the ones worth surfacing)
        ddpm_T: int = 200,
        ddpm_hidden_dim: int = 128,
        ddpm_t_embed_dim: int = 64,
        ddpm_learning_rate: float = 3e-4,
        ddpm_n_iter: int = 200_000,
        ddpm_ema_decay: float = 0.999,
        ddpm_beta_max: float = 0.02,
        ddpm_batch_size: int = 256,
        ddpm_verbose_every: int = 0,
        ddpm_eps: float = 1e-5,
        # runtime
        ddpm_device: str = "auto",
        key: Optional[jax.Array] = None,
        # ann
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        # “escape hatches”
        dmap_kwargs: Optional[Dict[str, Any]] = None,
        gplm_kwargs: Optional[Dict[str, Any]] = None,
        ddpm_kwargs: Optional[Dict[str, Any]] = None,
        # unicode aliases (β)
        **kwargs: Any,
    ):
        # ---- unicode aliases ----
        if "β" in kwargs:
            beta = float(kwargs.pop("β"))
        if "β_dmap" in kwargs:
            dmap_beta = float(kwargs.pop("β_dmap"))
        if "β_gplm" in kwargs:
            gplm_beta = float(kwargs.pop("β_gplm"))
        if kwargs:
            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")

        self.config = DIMAConfig(d=int(d), beta=float(beta), ddpm_device=str(ddpm_device))

        # devices + rng
        self.ddpm_device = _select_device(ddpm_device)
        self.cpu_device = _select_device("cpu")
        self.rng = random.PRNGKey(0) if key is None else key

        # training data
        self.R_iX = np.asarray(R_iX)
        if self.R_iX.ndim != 2:
            raise ValueError("R_iX must be 2D (N,D).")
        self.N, self.D = self.R_iX.shape
        self.d = int(d)

        # betas
        self.beta = float(beta)
        self.β = self.beta

        self.dmap_beta = float(self.beta if dmap_beta is None else dmap_beta)
        self.gplm_beta = float(self.beta if gplm_beta is None else gplm_beta)

        t0 = time.time()

        # -------------------------
        # 1) Train DMAP (CPU)
        # -------------------------
        dmap_init = dict(
            d=self.d,
            beta=self.dmap_beta,
            alpha=float(dmap_alpha),
            t=float(dmap_t),
            k=dmap_k,
            ann_backend=ann_backend,
            ann_params=ann_params,
            n_jobs=n_jobs,
        )
        if dmap_kwargs:
            dmap_init.update(dict(dmap_kwargs))
        # enforce headline knobs
        dmap_init["d"] = self.d
        dmap_init["beta"] = self.dmap_beta
        dmap_init["alpha"] = float(dmap_alpha)
        dmap_init["t"] = float(dmap_t)
        dmap_init["k"] = dmap_k

        self.enc = DMAP(self.R_iX, **dmap_init)

        # raw DMAP coordinates for *all* training points (Nyström OOS on training set)
        R_ix = np.asarray(self.enc(self.R_iX), dtype=np.float64)  # (N,d)

        # -------------------------
        # 2) Latent normalization for DDPM
        # -------------------------
        self.lat_mean_np = R_ix.mean(axis=0)
        self.lat_std_np = np.maximum(R_ix.std(axis=0), 1e-12)

        self.lat_mean_j = jax.device_put(jnp.asarray(self.lat_mean_np, dtype=jnp.float32), self.ddpm_device)
        self.lat_std_j = jax.device_put(jnp.asarray(self.lat_std_np, dtype=jnp.float32), self.ddpm_device)

        Z_ix = (R_ix - self.lat_mean_np) / self.lat_std_np  # (N,d)

        # -------------------------
        # 3) Train GPLM (CPU)
        # -------------------------
        gplm_init = dict(
            beta=self.gplm_beta,
            m=int(gplm_m),
            pred_k=gplm_pred_k,
            ann_backend=ann_backend,
            ann_params=ann_params,
            n_jobs=n_jobs,
        )
        if gplm_kwargs:
            gplm_init.update(dict(gplm_kwargs))
        # enforce headline knobs
        gplm_init["beta"] = self.gplm_beta
        gplm_init["m"] = int(gplm_m)
        gplm_init["pred_k"] = gplm_pred_k

        self.dec = GPLM(R_ix.astype(np.float32, copy=False), self.R_iX, **gplm_init)

        # -------------------------
        # 4) Train DDPM on normalized latents (DDPM device)
        # -------------------------
        Z_ix_j = jax.device_put(jnp.asarray(Z_ix, dtype=jnp.float32), self.ddpm_device)

        ddpm_init = dict(
            T=int(ddpm_T),
            hidden_dim=int(ddpm_hidden_dim),
            t_embed_dim=int(ddpm_t_embed_dim),
            learning_rate=float(ddpm_learning_rate),
            n_iter=int(ddpm_n_iter),
            ema_decay=float(ddpm_ema_decay),
            beta_max=float(ddpm_beta_max),
            batch_size=int(ddpm_batch_size),
            key=self.rng,
            verbose_every=int(ddpm_verbose_every),
            eps=float(ddpm_eps),
        )
        if ddpm_kwargs:
            ddpm_init.update(dict(ddpm_kwargs))
        # enforce headline knobs
        ddpm_init["T"] = int(ddpm_T)
        ddpm_init["hidden_dim"] = int(ddpm_hidden_dim)
        ddpm_init["t_embed_dim"] = int(ddpm_t_embed_dim)
        ddpm_init["learning_rate"] = float(ddpm_learning_rate)
        ddpm_init["n_iter"] = int(ddpm_n_iter)
        ddpm_init["ema_decay"] = float(ddpm_ema_decay)
        ddpm_init["beta_max"] = float(ddpm_beta_max)
        ddpm_init["batch_size"] = int(ddpm_batch_size)
        ddpm_init["verbose_every"] = int(ddpm_verbose_every)
        ddpm_init["eps"] = float(ddpm_eps)

        with jax.default_device(self.ddpm_device):
            self.dm = DDPM(Z_ix_j, **ddpm_init)

        self.training_time = time.time() - t0

    # -------------------------
    # Latent conversions
    # -------------------------
    def normalize(self, R_ax: Union[np.ndarray, jnp.ndarray]) -> jnp.ndarray:
        """raw latents (R) -> normalized latents (Z) on ddpm_device."""
        R = np.asarray(R_ax, dtype=np.float64)
        if R.ndim == 1:
            R = R[None, :]
        Z = (R - self.lat_mean_np) / self.lat_std_np
        Zj = jnp.asarray(Z, dtype=jnp.float32)
        return jax.device_put(Zj, self.ddpm_device)

    def unnormalize(self, Z_ax: Union[np.ndarray, jnp.ndarray]) -> np.ndarray:
        """normalized latents (Z) -> raw latents (R) on CPU (np)."""
        if isinstance(Z_ax, jax.Array):
            Z_np = np.asarray(jax.device_get(Z_ax))
        else:
            Z_np = np.asarray(Z_ax)
        if Z_np.ndim == 1:
            Z_np = Z_np[None, :]
        R = Z_np * self.lat_std_np + self.lat_mean_np
        return np.asarray(R)

    # -------------------------
    # Encode / Decode (public: raw latents)
    # -------------------------
    def encode(self, R_aX: Union[np.ndarray, jnp.ndarray], *, normalize: bool = False) -> Union[np.ndarray, jnp.ndarray]:
        """
        ambient -> raw DMAP latents (np) by default.
        If normalize=True, returns normalized latents (jnp) on ddpm_device.
        """
        X = np.asarray(R_aX)
        R_raw = np.asarray(self.enc(X))  # CPU, (a,d)
        if not normalize:
            return R_raw
        return self.normalize(R_raw)

    def decode(
        self,
        R_ax: Union[np.ndarray, jnp.ndarray],
        *,
        refine: bool = False,
        t_start: int = 10,
        add_noise: bool = True,
        key: Optional[jax.Array] = None,
        batch_size: Optional[int] = None,
    ) -> np.ndarray:
        """
        raw latent -> (optional DDPM refine in normalized coords) -> raw latent -> ambient.
        Returns ambient np.ndarray on CPU.
        """
        R_raw = np.asarray(R_ax, dtype=np.float64)
        single = (R_raw.ndim == 1)
        if single:
            R_raw = R_raw[None, :]

        if refine:
            Z = self.normalize(R_raw)  # on device
            Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
            R_raw = self.unnormalize(Z)  # back to CPU raw

        X_hat = self.dec(R_raw.astype(np.float32, copy=False), batch_size=batch_size)
        X_hat = np.asarray(X_hat)
        return X_hat[0] if single else X_hat

    def reconstruct(
        self,
        R_aX: Union[np.ndarray, jnp.ndarray],
        *,
        refine: bool = False,
        t_start: int = 10,
        add_noise: bool = True,
        key: Optional[jax.Array] = None,
        batch_size: Optional[int] = None,
    ) -> np.ndarray:
        """decode(encode(X))."""
        R_raw = self.encode(R_aX, normalize=False)
        return self.decode(R_raw, refine=refine, t_start=t_start, add_noise=add_noise, key=key, batch_size=batch_size)

    def sample(
        self,
        n: int,
        *,
        decode: bool = True,
        batch_size: Optional[int] = None,
    ) -> Union[np.ndarray, np.ndarray]:
        """
        Unconditional samples from latent DDPM.
        If decode=True: returns ambient samples (np) on CPU.
        If decode=False: returns raw latents (np) on CPU.
        """
        with jax.default_device(self.ddpm_device):
            Z = self.dm.sample(int(n))
        R = self.unnormalize(Z)  # raw (np)
        if not decode:
            return R
        return self.dec(R.astype(np.float32, copy=False), batch_size=batch_size)

    # -------------------------
    # Geodesic-ish flow wrapper (delegates to GPLM.flow)
    # -------------------------
    def flow(
        self,
        R_ax: Union[np.ndarray, jnp.ndarray],
        v_ax: Union[np.ndarray, jnp.ndarray],
        *,
        dt: float = 0.05,
        reg: float = 1e-8,
        keep_speed: bool = True,
        # optional DDPM projection step (in normalized coords)
        refine: bool = False,
        t_start: int = 10,
        add_noise: bool = True,
        key: Optional[jax.Array] = None,
        # decode return
        decode: bool = False,
        batch_size: Optional[int] = None,
    ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]]:
        """
        One step of latent flow in *raw* coordinates.
        Requires: your GPLM class implements:
            R_next, v_next = gplm.flow(R, v, dt=..., reg=..., keep_speed=...)
        If refine=True, we project the *position* through DDPM in normalized coords after the step.
        (Velocity after projection is left unchanged—projection isn’t a deterministic diffeo.)
        Returns:
          if decode=False:
            (R_next, v_next)        both np arrays
          if decode=True:
            (X_next, R_next, v_next)
        """
        R = np.asarray(R_ax, dtype=np.float64)
        v = np.asarray(v_ax, dtype=np.float64)
        single = (R.ndim == 1)
        if single:
            R = R[None, :]
            v = v[None, :]

        if not hasattr(self.dec, "flow"):
            raise AttributeError(
                "Decoder does not have .flow(). Make sure you updated GPLM to include flow()."
            )

        Rn, vn = self.dec.flow(R, v, dt=float(dt), reg=float(reg), keep_speed=bool(keep_speed))

        if refine:
            Z = self.normalize(Rn)  # device
            Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
            Rn = self.unnormalize(Z)

        if not decode:
            if single:
                return np.asarray(Rn[0]), np.asarray(vn[0])
            return np.asarray(Rn), np.asarray(vn)

        Xn = self.dec(Rn.astype(np.float32, copy=False), batch_size=batch_size)
        Xn = np.asarray(Xn)
        if single:
            return Xn[0], np.asarray(Rn[0]), np.asarray(vn[0])
        return Xn, np.asarray(Rn), np.asarray(vn)

    # -------------------------
    # Convenience __call__
    # -------------------------
    def __call__(
        self,
        A: Union[np.ndarray, jnp.ndarray],
        *,
        refine: bool = False,
        t_start: int = 10,
        add_noise: bool = True,
        key: Optional[jax.Array] = None,
        batch_size: Optional[int] = None,
        normalize_latent: bool = False,
    ) -> Union[np.ndarray, jnp.ndarray]:
        """
        Dispatch by last dimension:
          - if A is (a,D): encode -> raw latents (np) by default
          - if A is (a,d): decode -> ambient (np)
        Options:
          - normalize_latent=True only affects encoding (returns Z on device)
          - refine/t_start/add_noise/key only affect decoding
        """
        A_np = np.asarray(A)
        if A_np.ndim == 1:
            A_np = A_np[None, :]

        if A_np.shape[1] == self.D:
            return self.encode(A_np, normalize=bool(normalize_latent))

        if A_np.shape[1] == self.d:
            return self.decode(
                A_np,
                refine=bool(refine),
                t_start=int(t_start),
                add_noise=bool(add_noise),
                key=key,
                batch_size=batch_size,
            )

        raise ValueError(f"Input has last-dim {A_np.shape[1]}, expected D={self.D} or d={self.d}.")

    # -------------------------
    # Save / Load
    # -------------------------
    def _pack_encoder(self) -> Dict[str, Any]:
        enc = self.enc
    
        # qalpha_i (ascii/unicode)
        qalpha_i = np.asarray(getattr(enc, "qalpha_i", getattr(enc, "qα_i")))
    
        # 1) Try to read R_over_* from whichever name exists
        R_over = getattr(enc, "R_over_lam_ix", None)
        if R_over is None:
            R_over = getattr(enc, "R_over_λ_ix", None)
    
        # 2) If missing, compute it from R_ix and λ_x (most robust)
        if R_over is None:
            R_ix = getattr(enc, "R_ix", None)
    
            # λ_x is unicode in upstream; add fallbacks for safety
            lam = getattr(enc, "λ_x", None)
            if lam is None:
                lam = getattr(enc, "lam_x", None)
            if lam is None:
                lam = getattr(enc, "lambda_x", None)
    
            if R_ix is None or lam is None:
                raise AttributeError(
                    "DMAP encoder is missing R_over_{λ,lam}_ix and also lacks (R_ix, λ_x) "
                    "to reconstruct it. Please ensure your DMAP computes diffusion coords."
                )
    
            lam = np.asarray(lam, dtype=np.float64)
            lam = np.maximum(lam, 1e-30)  # avoid division by 0
            R_ix = np.asarray(R_ix, dtype=np.float64)
    
            R_over = (R_ix / lam[None, :]).astype(np.float64, copy=False)
    
        return dict(
            R_iX=np.asarray(enc.R_iX),
            qalpha_i=qalpha_i,
            # Always store with ASCII key expected by FrozenDMAP loader
            R_over_lam_ix=np.asarray(R_over, dtype=np.float64),
            k=int(enc.k),
            beta=float(getattr(enc, "beta", getattr(enc, "β"))),
            alpha=float(getattr(enc, "alpha", getattr(enc, "α"))),
            eps=float(getattr(enc, "eps", getattr(enc, "ε"))),
            dtype=_np_dtype_str(getattr(enc, "dtype", np.float32)),
        )
    
    
    def _pack_decoder(self) -> Dict[str, Any]:
        return dict(
            Z_mx_w=np.asarray(getattr(self.dec, "Z_mx_w", None)),
            M_mX=np.asarray(self.dec.M_mX),
            mean_X=np.asarray(getattr(self.dec, "mean_X", np.zeros((self.D,), dtype=np.float64))),
            lat_mean_x=np.asarray(getattr(self.dec, "lat_mean_x", np.zeros((self.d,), dtype=np.float64))),
            lat_std_x=np.asarray(getattr(self.dec, "lat_std_x", np.ones((self.d,), dtype=np.float64))),
            beta=float(getattr(self.dec, "beta", getattr(self.dec, "β"))),
            eps=float(getattr(self.dec, "eps", getattr(self.dec, "ε"))),
            pred_k=getattr(self.dec, "pred_k", getattr(self.dec, "pred_κ", None)),
            dtype=_np_dtype_str(getattr(self.dec, "dtype", np.float32)),
        )

    def state_dict(self) -> Dict[str, Any]:
        dd = dict(
            T=int(self.dm.T),
            D=int(self.dm.D),
            hidden_dim=int(getattr(self.dm.model, "hidden", 128)),
            t_embed_dim=int(getattr(self.dm.model, "t_dim", 64)),
            ema_decay=float(getattr(self.dm, "ema_decay", 0.999)),
            beta_max=float(getattr(self.dm, "beta_max", 0.02)),
            eps=float(getattr(self.dm, "eps", 1e-5)),
            params=self.dm.state.params,
            ema_params=self.dm.state.ema_params,
        )

        state = dict(
            meta=dict(
                N=int(self.N),
                D=int(self.D),
                d=int(self.d),
                training_time=float(getattr(self, "training_time", 0.0)),
            ),
            config=asdict(self.config),
            latent_norm=dict(
                mean=np.asarray(self.lat_mean_np, dtype=np.float64),
                std=np.asarray(self.lat_std_np, dtype=np.float64),
            ),
            encoder=self._pack_encoder(),
            decoder=self._pack_decoder(),
            ddpm=dd,
        )
        return state

    def save_local(self, weights_file: str = "dima.msgpack", config_file: str = "config.json") -> None:
        """
        Save full DIMA state to a msgpack file + a readable JSON config.

        This version is robust to accidental tuples inside the state tree
        (msgpack cannot serialize tuples by default).
        """
        def _sanitize(x):
            # Convert tuples -> lists recursively (msgpack-safe)
            if isinstance(x, tuple):
                return [_sanitize(v) for v in x]
            if isinstance(x, list):
                return [_sanitize(v) for v in x]
            if isinstance(x, dict):
                return {k: _sanitize(v) for k, v in x.items()}
            return x

        state = _sanitize(self.state_dict())
        blob = flax_ser.msgpack_serialize(state)

        with open(weights_file, "wb") as f:
            f.write(blob)

        with open(config_file, "w") as f:
            json.dump(state["config"], f, indent=2)

        return None


    @classmethod
    def load_local(
        cls,
        weights_file: str = "dima.msgpack",
        *,
        ddpm_device: str = "auto",
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        key: Optional[jax.Array] = None,
    ) -> "DIMA":
        with open(weights_file, "rb") as f:
            state = flax_ser.msgpack_restore(f.read())

        obj = cls.__new__(cls)  # bypass __init__

        obj.config = DIMAConfig(**state["config"])
        obj.ddpm_device = _select_device(ddpm_device)
        obj.cpu_device = _select_device("cpu")
        obj.training_time = float(state["meta"].get("training_time", 0.0))

        obj.N = int(state["meta"]["N"])
        obj.D = int(state["meta"]["D"])
        obj.d = int(state["meta"]["d"])

        # RNG
        obj.rng = random.PRNGKey(0) if key is None else key

        # beta (for display / convenience)
        obj.beta = float(obj.config.beta)
        obj.β = obj.beta

        # latent norm
        obj.lat_mean_np = np.asarray(state["latent_norm"]["mean"], dtype=np.float64)
        obj.lat_std_np = np.asarray(state["latent_norm"]["std"], dtype=np.float64)

        obj.lat_mean_j = jax.device_put(jnp.asarray(obj.lat_mean_np, dtype=jnp.float32), obj.ddpm_device)
        obj.lat_std_j = jax.device_put(jnp.asarray(obj.lat_std_np, dtype=jnp.float32), obj.ddpm_device)

        # frozen encoder + rehydrated decoder-as-GPLM (so flow works if GPLM.flow exists)
        obj.enc = FrozenDMAP(state["encoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
        obj.dec = _restore_gplm_as_object(state["decoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)

        # rebuild DDPM skeleton with dummy data, then load params
        dd = state["ddpm"]
        T = int(dd["T"])
        D = int(dd["D"])
        hidden_dim = int(dd["hidden_dim"])
        t_embed_dim = int(dd["t_embed_dim"])
        ema_decay = float(dd.get("ema_decay", 0.999))
        beta_max = float(dd.get("beta_max", 0.02))
        eps = float(dd.get("eps", 1e-5))

        dummy = jnp.zeros((1, D), dtype=jnp.float32)
        with jax.default_device(obj.ddpm_device):
            obj.dm = DDPM(
                dummy,
                T=T,
                hidden_dim=hidden_dim,
                t_embed_dim=t_embed_dim,
                learning_rate=1e-3,
                n_iter=0,               # skip training on load
                ema_decay=ema_decay,
                beta_max=beta_max,
                batch_size=1,
                key=obj.rng,
                verbose_every=0,
                eps=eps,
            )
            obj.dm.state = obj.dm.state.replace(params=dd["params"], ema_params=dd["ema_params"])

        obj.R_iX = None  # training data not stored by default
        return obj

    # -------------------------
    # (Optional) HF helpers
    # -------------------------

    def upload_to_huggingface(
        self,
        repo_id: str,
        *,
        weights_file: str = "dima.msgpack",
        config_file: str = "config.json",
        token: Optional[str] = None,
        repo_type: str = "model",
        revision: Optional[str] = None,
    ) -> None:
        """
        Serialize locally (weights + config) and upload them to Hugging Face Hub.

        Parameters
        ----------
        repo_id : str
            e.g. "username/my-dima-model"
        weights_file : str
            Filename used both locally and in the HF repo.
        config_file : str
            Human-readable JSON config filename (also uploaded).
        token : Optional[str]
            HF token. If None, uses HfFolder.get_token().
        repo_type : str
            Usually "model".
        revision : Optional[str]
            Optional target branch/revision (if your hub client supports it).
        """
        if not _HAS_HF:
            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).")

        # 1) Save artifacts locally
        self.save_local(weights_file=weights_file, config_file=config_file)

        # 2) Resolve token
        if token is None:
            token = HfFolder.get_token()
        if token is None:
            raise RuntimeError("No HF token found. Provide `token=...` or run `huggingface-cli login`.")

        # 3) Create repo if needed
        api = HfApi()
        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)

        # 4) Upload files
        common_kwargs = dict(repo_id=repo_id, repo_type=repo_type, token=token)
        if revision is not None:
            common_kwargs["revision"] = revision

        upload_file(
            path_or_fileobj=weights_file,
            path_in_repo=weights_file,
            **common_kwargs,
        )
        upload_file(
            path_or_fileobj=config_file,
            path_in_repo=config_file,
            **common_kwargs,
        )
        return None


    @classmethod
    def download_from_huggingface(
        cls,
        repo_id: str,
        *,
        weights_file: str = "dima.msgpack",
        ddpm_device: str = "auto",
        ann_backend: "ANNBackend" = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        key: Optional["jax.Array"] = None,
        token: Optional[str] = None,
        repo_type: str = "model",
        revision: Optional[str] = None,
    ) -> "DIMA":
        """
        Download weights from HF Hub and rehydrate a DIMA object via load_local.

        Returns
        -------
        DIMA
            A ready-to-use (inference) DIMA instance.
        """
        if not _HAS_HF:
            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).")

        if token is None:
            token = HfFolder.get_token()

        dl_kwargs = dict(repo_id=repo_id, filename=weights_file, repo_type=repo_type)
        if token is not None:
            dl_kwargs["token"] = token
        if revision is not None:
            dl_kwargs["revision"] = revision

        path = hf_hub_download(**dl_kwargs)

        return cls.load_local(
            path,
            ddpm_device=ddpm_device,
            ann_backend=ann_backend,
            ann_params=ann_params,
            n_jobs=n_jobs,
            key=key,
        )


    # Backward-compatible aliases (the upstream file uses save_hf/load_hf)
    def save_hf(self, repo_id: str, weights_file: str = "dima.msgpack", config_file: str = "config.json") -> None:
        return self.upload_to_huggingface(repo_id, weights_file=weights_file, config_file=config_file)

    @classmethod
    def load_hf(
        cls,
        repo_id: str,
        *,
        weights_file: str = "dima.msgpack",
        ddpm_device: str = "auto",
        ann_backend: "ANNBackend" = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        key: Optional["jax.Array"] = None,
    ) -> "DIMA":
        return cls.download_from_huggingface(
            repo_id,
            weights_file=weights_file,
            ddpm_device=ddpm_device,
            ann_backend=ann_backend,
            ann_params=ann_params,
            n_jobs=n_jobs,
            key=key,
        )

__all__ = ["DIMA", "DIMAConfig"]