# --- src/dima/ddpm.py
# +++ src/dima/ddpm.py
# @@ -368,4 +368,204 @@
#          return self.reverse_from_T(noise)
# src/dima/ddpm.py
from __future__ import annotations

from typing import Any, Optional

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

from flax import linen as nn
from flax.training import train_state
from flax import struct
import optax
from .ann import ANNBackend, make_ann
from .ddpm import DDPM
# +from .ddpm import DDPM, DDIM
from .dmap import DMAP
from .gplm import GPLM

__all__ = ["DDPM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]


# ------------------------------
# DDIM (deterministic sampler; eta=0.0)
# ------------------------------

class DDIM(DDPM):
    """
    DDIM sampler for D-dimensional latents.

    Notes
    -----
    * Training is identical to DDPM (epsilon prediction).
    * Sampling / refinement uses the DDIM update:
        x_{t_prev} = sqrt(a_bar_prev) * x0_hat
                   + sqrt(1 - a_bar_prev - sigma^2) * eps_pred
                   + sigma * z,
      where sigma = eta * sqrt( (1-a_bar_prev)/(1-a_bar_t) * (1 - a_bar_t/a_bar_prev) ).

    With eta=0.0 the reverse process is deterministic conditioned on x_t.

    Added (HF I/O):
      - state_dict / save_local / load_local / from_state
      - upload_to_huggingface / download_from_huggingface
    """

    def __init__(
        self,
        Z_iX: jnp.ndarray,
        *,
        eta: float = 0.0,
        steps: Optional[int] = None,
        **ddpm_kwargs: Any,
    ):
        super().__init__(Z_iX, **ddpm_kwargs)
        self.eta = float(eta)
        self.steps = None if steps is None else int(steps)

    @staticmethod
    def _make_ddim_step(params_ema, apply_fn, alpha_bar_s, eps: float, eta: float):
        eta = float(eta)

        @jax.jit
        def step(carry, t_pair):
            key, x = carry  # x: (B, D)
            t, t_prev = t_pair  # scalar int32

            # gather ā_t and ā_prev (supports skipping)
            a_bar_t = jnp.clip(alpha_bar_s[t], eps, 1.0)
            a_bar_prev = jnp.clip(alpha_bar_s[t_prev], eps, 1.0)

            B = x.shape[0]
            t_batch = jnp.full((B, 1), t, dtype=jnp.int32)

            eps_pred = apply_fn({"params": params_ema}, x, t_batch)  # (B, D)

            # x0 estimate
            sqrt_one_minus = jnp.sqrt(jnp.clip(1.0 - a_bar_t, eps, 1.0))
            x0_hat = (x - sqrt_one_minus * eps_pred) / jnp.sqrt(a_bar_t)

            # sigma_t (controls stochasticity). eta=0 => sigma=0 (deterministic)
            frac = (1.0 - a_bar_prev) / jnp.clip(1.0 - a_bar_t, eps, 1.0)
            inside = 1.0 - a_bar_t / jnp.clip(a_bar_prev, eps, 1.0)
            sigma = eta * jnp.sqrt(jnp.clip(frac * inside, 0.0, 1.0))

            # direction term
            dir_coeff = jnp.sqrt(jnp.clip(1.0 - a_bar_prev - sigma**2, 0.0, 1.0))
            x_prev = jnp.sqrt(a_bar_prev) * x0_hat + dir_coeff * eps_pred

            # optional noise injection (eta>0)
            if eta > 0.0:
                key, k = random.split(key)
                z = random.normal(k, x.shape)
                x_prev = x_prev + sigma * z

            return (key, x_prev), x_prev

        return step

    def _t_schedule(self, t_start: int, *, steps: Optional[int] = None) -> jnp.ndarray:
        """
        Descending integer timesteps from t_start -> 0.

        If steps is provided (or self.steps), we use a reduced DDIM schedule
        by rounding a linear grid and making it unique.
        """
        t_start = int(t_start)
        if t_start <= 0:
            return jnp.array([0], dtype=jnp.int32)

        S = self.steps if steps is None else int(steps)
        if S is None or S >= (t_start + 1):
            return jnp.arange(t_start, -1, -1, dtype=jnp.int32)

        grid = jnp.linspace(float(t_start), 0.0, int(S), dtype=jnp.float32)
        ts = jnp.unique(jnp.round(grid).astype(jnp.int32))
        ts = ts[::-1]  # descending

        # guarantee endpoints
        if ts[0] != t_start:
            ts = jnp.concatenate([jnp.array([t_start], dtype=jnp.int32), ts])
        if ts[-1] != 0:
            ts = jnp.concatenate([ts, jnp.array([0], dtype=jnp.int32)])
        return ts

    def refine_latents(
        self,
        z0: jnp.ndarray,
        t_start: int = 10,
        key: Optional[jax.Array] = None,
        add_noise: bool = True,
        *,
        eta: Optional[float] = None,
        steps: Optional[int] = None,
    ) -> jnp.ndarray:
        """
        Refine latents by:
          (optional) forward-noise z0 to step t_start
          DDIM reverse from t_start -> 0 using EMA params.
        """
        z0 = jnp.asarray(z0, dtype=jnp.float32)
        if z0.ndim != 2 or z0.shape[1] != self.D:
            raise ValueError(f"z0 must have shape (B,{self.D}).")
        if not (0 <= int(t_start) < self.T):
            raise ValueError(f"t_start must be in [0, {self.T-1}]")
        t_start = int(t_start)

        # RNG
        if key is None:
            self.key, key = random.split(self.key)
        else:
            # advance internal RNG too
            self.key, _ = random.split(key)

        # optional forward noise
        if add_noise and t_start > 0:
            key, k_eps = random.split(key)
            eps_noise = random.normal(k_eps, z0.shape)
            a_bar_t = jnp.clip(self.alpha_bar_s[t_start], self.eps, 1.0) # Corrected attribute name
            z_t = jnp.sqrt(a_bar_t) * z0 + jnp.sqrt(1.0 - a_bar_t) * eps_noise
        else:
            z_t = z0

        # schedule
        ts = self._t_schedule(t_start, steps=steps)
        if ts.shape[0] == 1:
            return z_t

        t_pairs = jnp.stack([ts[:-1], ts[1:]], axis=1)  # (K,2)

        step = self._make_ddim_step(
            self.state.ema_params,
            self.state.apply_fn,
            self.alpha_bar_s, # Corrected attribute name
            self.eps,
            self.eta if eta is None else float(eta),
        )

        (final_key, _), trace = jax.lax.scan(step, (key, z_t), xs=t_pairs)

        self.key = final_key
        return trace[-1]

    def __call__(
        self,
        z0: jnp.ndarray,
        t_start: int = 10,
        key: Optional[jax.Array] = None,
        add_noise: bool = True,
        *,
        eta: Optional[float] = None,
        steps: Optional[int] = None,
    ) -> jnp.ndarray:
        return self.refine_latents(
            z0, t_start=t_start, key=key, add_noise=add_noise, eta=eta, steps=steps
        )

    def reverse_from_T(self, x_T: jnp.ndarray, *, eta: Optional[float] = None, steps: Optional[int] = None) -> jnp.ndarray:
        x_T = jnp.asarray(x_T, dtype=jnp.float32)
        if x_T.ndim != 2 or x_T.shape[1] != self.D:
            raise ValueError(f"x_T must have shape (B,{self.D}).")

        self.key, k0 = random.split(self.key)

        ts = self._t_schedule(self.T - 1, steps=steps)
        t_pairs = jnp.stack([ts[:-1], ts[1:]], axis=1)

        step = self._make_ddim_step(
            self.state.ema_params,
            self.state.apply_fn,
            self.alpha_bar_s, # Corrected attribute name
            self.eps,
            self.eta if eta is None else float(eta),
        )

        (_, _), trace = jax.lax.scan(step, (k0, x_T), xs=t_pairs)
        return trace[-1]

    def sample(self, N: int = 10_000, *, eta: Optional[float] = None, steps: Optional[int] = None) -> jnp.ndarray:
        self.key, k = random.split(self.key)
        noise = random.normal(k, (int(N), self.D))
        return self.reverse_from_T(noise, eta=eta, steps=steps)

    # ------------------------------------------------------------------
    # Serialization + Hugging Face Hub I/O (integrated)
    # ------------------------------------------------------------------

    def state_dict(self):
        """
        Msgpack-safe checkpoint dict (includes DDIM-specific eta/steps).
        Uses flax.serialization.to_state_dict for TrainStateEMA to avoid tuple issues (opt_state).
        """
        from flax import serialization as flax_ser  # type: ignore
        import numpy as np

        return dict(
            # DDPM core hyperparameters
            T=int(self.T),
            D=int(self.D),
            hidden_dim=int(getattr(self, "hidden_dim", 0) or self.hidden_dim),
            t_embed_dim=int(getattr(self, "t_embed_dim", 0) or self.t_embed_dim),
            learning_rate=float(self.learning_rate),
            ema_decay=float(self.ema_decay),
            beta_max=float(self.beta_max),
            batch_size=None if (self.batch_size is None) else int(self.batch_size),
            eps=float(self.eps),

            # DDIM specifics
            eta=float(self.eta),
            steps=None if (self.steps is None) else int(self.steps),

            # RNG key (portable)
            key=np.asarray(self.key),

            # TrainStateEMA (params, ema_params, opt_state, step)
            train_state=flax_ser.to_state_dict(self.state),
        )

    def save_local(self, weights_file: str = "ddim.msgpack", config_file: Optional[str] = "ddim_config.json") -> None:
        from flax import serialization as flax_ser  # type: ignore
        import json as _json

        ckpt = self.state_dict()
        blob = flax_ser.msgpack_serialize(ckpt)

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

        if config_file is not None:
            cfg = dict(
                T=int(ckpt["T"]),
                D=int(ckpt["D"]),
                hidden_dim=int(ckpt["hidden_dim"]),
                t_embed_dim=int(ckpt["t_embed_dim"]),
                learning_rate=float(ckpt["learning_rate"]),
                ema_decay=float(ckpt["ema_decay"]),
                beta_max=float(ckpt["beta_max"]),
                batch_size=ckpt["batch_size"],
                eps=float(ckpt["eps"]),
                eta=float(ckpt["eta"]),
                steps=ckpt["steps"],
            )
            with open(config_file, "w") as f:
                _json.dump(cfg, f, indent=2)

    @classmethod
    def from_state(cls, ckpt, *, key: Optional[jax.Array] = None) -> "DDIM":
        """
        Rehydrate a DDIM instance without retraining:
          - rebuilds a skeleton DDIM (which builds schedules/model/optimizer)
          - restores TrainStateEMA via flax.serialization.from_state_dict
        """
        from flax import serialization as flax_ser  # type: ignore
        import numpy as np

        T = int(ckpt["T"])
        D = int(ckpt["D"])
        hidden_dim = int(ckpt["hidden_dim"])
        t_embed_dim = int(ckpt["t_embed_dim"])
        learning_rate = float(ckpt["learning_rate"])
        ema_decay = float(ckpt["ema_decay"])
        beta_max = float(ckpt["beta_max"])
        batch_size = ckpt.get("batch_size", None)
        eps = float(ckpt.get("eps", 1e-5))

        eta = float(ckpt.get("eta", 0.0))
        steps = ckpt.get("steps", None)
        steps = None if (steps is None) else int(steps)

        if key is None:
            if "key" in ckpt:
                key = jnp.asarray(np.asarray(ckpt["key"]))
            else:
                key = random.PRNGKey(0)

        # skeleton (no training)
        dummy = jnp.zeros((1, D), dtype=jnp.float32)
        obj = cls(
            dummy,
            T=T,
            hidden_dim=hidden_dim,
            t_embed_dim=t_embed_dim,
            learning_rate=learning_rate,
            n_iter=0,
            ema_decay=ema_decay,
            beta_max=beta_max,
            batch_size=batch_size,
            key=key,
            verbose_every=0,
            eps=eps,
            eta=eta,
            steps=steps,
        )

        obj.state = flax_ser.from_state_dict(obj.state, ckpt["train_state"])

        # restore RNG key and DDIM attrs
        if "key" in ckpt:
            obj.key = jnp.asarray(np.asarray(ckpt["key"]))
        obj.eta = eta
        obj.steps = steps

        return obj

    @classmethod
    def load_local(cls, weights_file: str = "ddim.msgpack", *, key: Optional[jax.Array] = None) -> "DDIM":
        from flax import serialization as flax_ser  # type: ignore

        with open(weights_file, "rb") as f:
            ckpt = flax_ser.msgpack_restore(f.read())
        return cls.from_state(ckpt, key=key)

    def upload_to_huggingface(
        self,
        repo_id: str,
        *,
        weights_file: str = "ddim.msgpack",
        config_file: str = "ddim_config.json",
        token: Optional[str] = None,
        repo_type: str = "model",
        revision: Optional[str] = None,
    ) -> None:
        """
        Upload DDIM checkpoint to HF Hub (msgpack + JSON).
        Mirrors the DMAP/DDPM pattern.
        """
        try:
            from huggingface_hub import HfApi, HfFolder, upload_file  # type: ignore
        except Exception as e:
            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e

        self.save_local(weights_file=weights_file, config_file=config_file)

        if token is None:
            token = HfFolder.get_token()
        if token is None:
            raise RuntimeError("No HF token found. Run `huggingface-cli login`, or pass `token=...`.")

        import os as _os

        api = HfApi()
        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)

        wf = _os.path.basename(weights_file)
        cf = _os.path.basename(config_file)

        upload_file(
            path_or_fileobj=weights_file,
            path_in_repo=wf,
            repo_id=repo_id,
            token=token,
            repo_type=repo_type,
            revision=revision,
        )
        upload_file(
            path_or_fileobj=config_file,
            path_in_repo=cf,
            repo_id=repo_id,
            token=token,
            repo_type=repo_type,
            revision=revision,
        )

    @classmethod
    def download_from_huggingface(
        cls,
        repo_id: str,
        *,
        weights_file: str = "ddim.msgpack",
        token: Optional[str] = None,
        repo_type: str = "model",
        revision: Optional[str] = None,
        key: Optional[jax.Array] = None,
    ) -> "DDIM":
        """
        Download DDIM checkpoint from HF Hub and return a rehydrated DDIM instance.
        """
        try:
            from huggingface_hub import hf_hub_download  # type: ignore
        except Exception as e:
            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e

        path = hf_hub_download(
            repo_id=repo_id,
            filename=weights_file,
            token=token,
            repo_type=repo_type,
            revision=revision,
        )
        return cls.load_local(path, key=key)

    # optional alias for naming symmetry
    download_to_huggingface = download_from_huggingface


__all__ = ["DDPM", "DDIM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]

# --- src/dima/dima.py
# +++ src/dima/dima.py
# @@ -15,7 +15,7 @@
#  from flax import serialization as flax_ser
 

 
# @@ -374,7 +374,13 @@
#          ddpm_init["eps"] = float(ddpm_eps)
 
#          with jax.default_device(self.ddpm_device):
# -            self.dm = DDPM(Z_ix_j, **ddpm_init)
# +            sampler = str(ddpm_init.pop("sampler", "ddpm")).lower()
# +            if sampler == "ddim":
# +                eta = float(ddpm_init.pop("eta", 0.0))
# +                steps = ddpm_init.pop("steps", None)
# +                self.dm = DDIM(Z_ix_j, eta=eta, steps=steps, **ddpm_init)
# +            else:
# +                self.dm = DDPM(Z_ix_j, **ddpm_init)
 
#          self.training_time = time.time() - t0