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

from typing import Any, Dict, Literal, Optional, Tuple, Union

import numpy as np
import scipy.linalg as la
import json
import os
import tempfile

from .ann import ANNBackend, make_ann
from .utils import fps_indices, median_eps_from_knn_d2, sqdist_ab


InducingMode = Literal["random_subset", "fps", "kmeans_medoids", "given"]




def _kmeans2_safe(Z: np.ndarray, m: int, seed: int = 0) -> np.ndarray:
    """
    KMeans centers with a safe fallback.
    Uses scipy.cluster.vq.kmeans2 if available; otherwise samples points.
    """
    Z = np.asarray(Z)
    m = int(min(max(1, m), Z.shape[0]))
    try:
        from scipy.cluster.vq import kmeans2  # type: ignore

        C, _ = kmeans2(Z.astype(np.float64, copy=False), m, minit="points", seed=seed)
        return C.astype(Z.dtype, copy=False)
    except Exception:
        rng = np.random.default_rng(seed)
        idx = rng.choice(Z.shape[0], size=m, replace=False)
        return Z[idx]


def _to_builtin(x: Any) -> Any:
    """
    Convierte tipos numpy / tuplas / dicts anidados a tipos nativos serializables.
    - Para msgpack (flax.serialization) las ndarrays se dejan como ndarrays.
    - Para JSON, este helper se usa solo sobre objetos escalares/dicts (sin ndarrays).
    """
    if x is None:
        return None
    if isinstance(x, (bool, int, float, str)):
        return x
    if isinstance(x, (np.integer, np.floating)):
        return x.item()
    if isinstance(x, (tuple, list)):
        return [_to_builtin(v) for v in x]
    if isinstance(x, dict):
        return {str(k): _to_builtin(v) for k, v in x.items()}
    # fallback conservador
    return str(x)


class GPLM:
    """
    Inducing-point / Nyström GP (kernel ridge) decoder on latents.

    Entrenamiento:
      - Construye puntos inductores Z_mx_w (en latente blanqueado o no),
      - Estima eps (si no se provee) a partir de distancias kNN,
      - Resuelve M_mX por Cholesky (Nyström KRR/GP mean).

    Además:
      - decode() vía __call__.
      - flow() integra un paso tipo generalized-leapfrog (geodesic flow) sobre el pullback manifold.
      - NUEVO: save_local/load_local + upload_to_huggingface/download_from_huggingface.
    """

    def __init__(
        self,
        R_ix: np.ndarray,
        R_iX: np.ndarray,
        *,
        # ASCII
        beta: float = 1.0,
        # eps estimation
        eps: Optional[float] = None,
        k_eps: int = 256,
        eps_use_kth: bool = True,
        eps_mul: float = 1.0,
        # regularization
        sigma2: float = 1e-5,
        jitter: float = 1e-8,
        # inducing
        m: int = 1024,
        inducing: InducingMode = "kmeans_medoids",
        Z_mx: Optional[np.ndarray] = None,
        seed: int = 0,
        # preprocess
        center_X: bool = True,
        whiten_latent: bool = False,
        dtype: Any = np.float32,
        # compute/memory
        fit_block: int = 8192,
        # inference
        pred_k: Optional[int] = None,
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        # unicode aliases
        **kwargs: Any,
    ):
        # ---- map unicode kwargs -> ascii
        if "β" in kwargs:
            beta = kwargs.pop("β")
        if "ε" in kwargs:
            eps = kwargs.pop("ε")
        if "κ_eps" in kwargs:
            k_eps = kwargs.pop("κ_eps")
        if "ε_use_kth" in kwargs:
            eps_use_kth = kwargs.pop("ε_use_kth")
        if "ε_mul" in kwargs:
            eps_mul = kwargs.pop("ε_mul")
        if "σ2" in kwargs:
            sigma2 = kwargs.pop("σ2")
        if "pred_κ" in kwargs:
            pred_k = kwargs.pop("pred_κ")
        if kwargs:
            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")

        # ---- store params (and keep enough meta to reconstruct on load)
        self.beta = float(beta)
        self.β = self.beta

        self.sigma2 = float(sigma2)
        self.σ2 = self.sigma2

        self.jitter = float(jitter)
        self.seed = int(seed)

        self.dtype = dtype
        self.dtype_str = str(np.dtype(dtype))

        self.fit_block = int(fit_block)
        self.center_X = bool(center_X)
        self.whiten_latent = bool(whiten_latent)
        self.inducing = inducing

        self.ann_backend = ann_backend
        self.ann_params = ann_params if ann_params is None else dict(ann_params)
        self.n_jobs = int(n_jobs)

        # ---- validate / cast
        R_ix = np.ascontiguousarray(np.asarray(R_ix).astype(self.dtype, copy=False))
        R_iX = np.ascontiguousarray(np.asarray(R_iX).astype(self.dtype, copy=False))
        if R_ix.ndim != 2 or R_iX.ndim != 2 or R_ix.shape[0] != R_iX.shape[0]:
            raise ValueError("R_ix must be (N,d) and R_iX must be (N,D) with same N.")

        self.R_ix = R_ix
        self.R_iX = R_iX
        self.N, self.d_lat = R_ix.shape
        _, self.D = R_iX.shape

        # ---- center output
        if self.center_X:
            self.mean_X = R_iX.mean(axis=0).astype(np.float64)
            Y = (R_iX.astype(np.float64) - self.mean_X[None, :])
        else:
            self.mean_X = np.zeros((self.D,), dtype=np.float64)
            Y = R_iX.astype(np.float64)

        # ---- latent whitening (optional)
        Ztrain = R_ix.astype(np.float64)
        if self.whiten_latent:
            self.lat_mean_x = Ztrain.mean(axis=0)
            self.lat_std_x = np.maximum(Ztrain.std(axis=0), 1e-12)
            Ztrain_w = (Ztrain - self.lat_mean_x) / self.lat_std_x
        else:
            self.lat_mean_x = np.zeros((self.d_lat,), dtype=np.float64)
            self.lat_std_x = np.ones((self.d_lat,), dtype=np.float64)
            Ztrain_w = Ztrain
        self.R_ix_w = Ztrain_w  # (N,d) float64

        # ---- ANN on training latents (eps + medoids snapping)
        self.ann_train, _ = make_ann(self.ann_backend, ann_params=self.ann_params, n_jobs=self.n_jobs)
        self.ann_train.build(self.R_ix_w.astype(self.dtype, copy=False))

        # ---- eps via kNN distances on latents
        if eps is None:
            k_eps = int(min(max(8, int(k_eps)), self.N - 1))
            j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), k_eps + 1)

            i = np.arange(self.N)[:, None]
            is_self = (j_iK1 == i)

            if np.any(is_self):
                D2_iK = np.empty((self.N, k_eps), dtype=np.float64)
                for ii in range(self.N):
                    keep = (j_iK1[ii] != ii)
                    D2_iK[ii] = D2_iK1[ii][keep][:k_eps]
            else:
                D2_iK = D2_iK1[:, :k_eps].astype(np.float64, copy=False)

            eps_hat = median_eps_from_knn_d2(D2_iK, use_kth=bool(eps_use_kth))
        else:
            eps_hat = float(eps)

        eps_hat *= float(eps_mul)
        if eps_hat <= 0:
            raise ValueError("eps must be > 0.")
        self.eps = float(eps_hat)
        self.ε = self.eps

        # ---- choose inducing points (in whitened latent space)
        rng = np.random.default_rng(self.seed)
        m_eff = int(min(max(1, int(m)), self.N))

        if Z_mx is not None:
            Zm = np.asarray(Z_mx, dtype=np.float64)
            if Zm.ndim != 2 or Zm.shape[1] != self.d_lat:
                raise ValueError("Z_mx must be (m, d_lat).")
            Zm_w = (Zm - self.lat_mean_x) / self.lat_std_x
        else:
            if inducing == "random_subset":
                idx = rng.choice(self.N, size=m_eff, replace=False)
                Zm_w = self.R_ix_w[idx]
            elif inducing == "fps":
                idx = fps_indices(self.R_ix_w, m=m_eff, seed=self.seed) # fps_indices is not defined
                # For now, fallback to random_subset if fps_indices is not available
                # idx = rng.choice(self.N, size=m_eff, replace=False)
                Zm_w = self.R_ix_w[idx]
            elif inducing == "kmeans_medoids":
                C = _kmeans2_safe(self.R_ix_w, m_eff, seed=self.seed).astype(np.float64, copy=False)
                j_cm, _ = self.ann_train.search(C.astype(self.dtype, copy=False), 1)
                idx = j_cm.reshape(-1).astype(np.int64)

                # de-duplicate and refill if needed
                idx_u = np.unique(idx)
                if idx_u.size < m_eff:
                    needed = m_eff - idx_u.size
                    pool = np.setdiff1d(np.arange(self.N), idx_u, assume_unique=False)
                    extra = rng.choice(pool, size=needed, replace=False) if pool.size >= needed else rng.choice(self.N, size=needed, replace=True)
                    idx = np.concatenate([idx_u, extra])
                else:
                    idx = idx_u[:m_eff]
                Zm_w = self.R_ix_w[idx]
            elif inducing == "given":
                raise ValueError("Provide Z_mx when inducing='given'.")
            else:
                raise ValueError(f"Unknown inducing mode: {inducing!r}")

        self.Z_mx_w = np.ascontiguousarray(Zm_w.astype(np.float64, copy=False))
        self.m = int(self.Z_mx_w.shape[0])

        # store raw inducing points (unwhitened) for convenience
        self.Z_mx = (self.Z_mx_w * self.lat_std_x[None, :]) + self.lat_mean_x[None, :]

        # ---- ANN on inducing points for fast prediction
        self.ann_Z, _ = make_ann(self.ann_backend, ann_params=self.ann_params, n_jobs=self.n_jobs)
        self.ann_Z.build(self.Z_mx_w.astype(self.dtype, copy=False))

        # pred_k
        if pred_k is None:
            self.pred_k = None
        else:
            self.pred_k = int(min(max(1, int(pred_k)), self.m))
        self.pred_κ = self.pred_k  # unicode alias

        # ---- W_mm and reduced solve
        D2_mm = sqdist_ab(self.Z_mx_w, self.Z_mx_w)
        W_mm = np.exp(-self.beta * (D2_mm.astype(np.float64) / self.eps))
        W_mm.flat[:: self.m + 1] += self.jitter
        self.W_mm = W_mm  # (m,m)

        G_mm = np.zeros((self.m, self.m), dtype=np.float64)
        B_mX = np.zeros((self.m, self.D), dtype=np.float64)
        bs = int(self.fit_block)

        for i0 in range(0, self.N, bs):
            i1 = min(self.N, i0 + bs)
            Zi = self.R_ix_w[i0:i1]  # (b,d)
            D2_im = sqdist_ab(Zi, self.Z_mx_w)
            C_im = np.exp(-self.beta * (D2_im.astype(np.float64) / self.eps))
            G_mm += C_im.T @ C_im
            B_mX += C_im.T @ Y[i0:i1]

        A_mm = G_mm + self.sigma2 * W_mm
        A_mm.flat[:: self.m + 1] += self.jitter

        cF = la.cho_factor(A_mm, lower=True, check_finite=False)
        self.M_mX = la.cho_solve(cF, B_mX, check_finite=False)  # (m,D)

    # ---------------------------
    # Inference
    # ---------------------------
    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:
            Y = self._decode(R_ax)
        else:
            bs = int(batch_size)
            out = []
            for s in range(0, R_ax.shape[0], bs):
                out.append(self._decode(R_ax[s : s + bs]))
            Y = np.vstack(out)

        return Y[0] if single else Y

    def _decode(self, R_ax: np.ndarray) -> np.ndarray:
        Za = R_ax.astype(np.float64, copy=False)
        Za_w = (Za - self.lat_mean_x) / self.lat_std_x

        if self.pred_k is None or self.pred_k == self.m:
            D2_am = sqdist_ab(Za_w, self.Z_mx_w)
            C_am = np.exp(-self.beta * (D2_am.astype(np.float64) / self.eps))
            Y = C_am @ self.M_mX
        else:
            j_aK, D2_aK = self.ann_Z.search(Za_w.astype(self.dtype, copy=False), self.pred_k)
            W = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)
            M = self.M_mX[j_aK]  # (a,k,D)
            Y = np.sum(W[:, :, None] * M, axis=1)  # (a,D)

        return Y + self.mean_X[None, :]

    # ============================================================
    # Geodesic flow on pullback manifold (no C_mm storage)
    # ============================================================
    def _rbf_cache_single(self, r_x: np.ndarray, *, idx_m: Optional[np.ndarray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Cache kernel terms at position r (single point).
        """
        c = self.beta / self.eps
        inv_std = 1.0 / self.lat_std_x

        r = r_x.astype(np.float64, copy=False)
        rw = (r - self.lat_mean_x) / self.lat_std_x

        Zw = self.Z_mx_w if idx_m is None else self.Z_mx_w[idx_m]
        Dw = rw[None, :] - Zw
        D2 = np.sum(Dw * Dw, axis=1)

        k_m = np.exp(-c * D2)  # (k,)
        Dw_over = Dw * inv_std[None, :]  # (k,d) = (r-z)/std^2
        grad_k = -(2.0 * c) * (k_m[:, None] * Dw_over)  # (k,d)

        return k_m, Dw_over, grad_k

    def _metric_from_gradM(
        self,
        grad_k: np.ndarray,  # (k,d)
        M_kX: np.ndarray,    # (k,D)
        *,
        D_block: int = 8192,
        lam: float = 1e-10,
    ) -> Tuple[np.ndarray, Tuple[np.ndarray, bool]]:
        """
        Compute pullback metric g = J^T J without forming C_mm.
        """
        k, d = grad_k.shape
        D = M_kX.shape[1]

        g = np.zeros((d, d), dtype=np.float64)
        for j0 in range(0, D, int(D_block)):
            j1 = min(D, j0 + int(D_block))
            Mb = M_kX[:, j0:j1]     # (k,block)
            Jb = Mb.T @ grad_k      # (block,d)
            g += Jb.T @ Jb          # (d,d)

        g = 0.5 * (g + g.T)
        g.flat[:: d + 1] += float(lam)

        cF = la.cho_factor(g, lower=True, check_finite=False)
        return g, cF

    def _force_from_cache_noC(
        self,
        k_m: np.ndarray,       # (k,)
        Dw_over: np.ndarray,   # (k,d)
        v_x: np.ndarray,       # (d,)
        M_kX: np.ndarray,      # (k,D)
        *,
        D_block: int = 8192,
    ) -> np.ndarray:
        """
        Geodesic momentum force without storing C_mm, using blocked contractions over ambient dim D.
        """
        c = self.beta / self.eps
        v = v_x.astype(np.float64, copy=False)

        inv_std2 = (1.0 / self.lat_std_x) ** 2  # (d,)
        s_m = Dw_over @ v                       # (k,)
        S = -(2.0 * c) * (k_m * s_m)            # (k,)

        term1 = (4.0 * c * c) * (k_m * s_m)[:, None] * Dw_over
        term2 = (2.0 * c) * k_m[:, None] * (v[None, :] * inv_std2[None, :])
        T = (term1 - term2).T                   # (d,k)

        d = v.shape[0]
        D = M_kX.shape[1]
        f = np.zeros((d,), dtype=np.float64)

        for j0 in range(0, D, int(D_block)):
            j1 = min(D, j0 + int(D_block))
            Mb = M_kX[:, j0:j1]   # (k,block)
            Jv_b = S @ Mb         # (block,)
            Hv_b = T @ Mb         # (d,block)
            f += Hv_b @ Jv_b      # (d,)

        return f

    def flow(
        self,
        R_ax: Union[np.ndarray, list],
        v_ax: Union[np.ndarray, list],
        *,
        eps: float = 1e-2,
        K_p: int = 5,
        K_q: int = 5,
        D_block: int = 8192,
        lam: float = 1e-10,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        One generalized-leapfrog step for geodesic flow on the pullback manifold.
        """
        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 R.ndim != 2 or v.ndim != 2 or R.shape != v.shape or R.shape[1] != self.d_lat:
            raise ValueError(f"Expected R_ax and v_ax shape (A,{self.d_lat}) (or ({self.d_lat},)).")

        A, _d = R.shape

        # Optional inducing subset indices per point
        if self.pred_k is None or self.pred_k == self.m:
            idx_aK = None
        else:
            Rw = (R - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
            idx_aK, _ = self.ann_Z.search(Rw.astype(self.dtype, copy=False), self.pred_k)
            idx_aK = idx_aK.astype(np.int64, copy=False)

        R_next = np.empty_like(R)
        v_next = np.empty_like(v)

        for a in range(A):
            idx = None if idx_aK is None else idx_aK[a]
            M_kX = self.M_mX if idx is None else self.M_mX[idx]

            r_n = R[a]
            v_n = v[a]

            # geometry at r_n
            k_m_n, Dw_over_n, grad_k_n = self._rbf_cache_single(r_n, idx_m=idx)
            g_n, cF_n = self._metric_from_gradM(grad_k_n, M_kX, D_block=D_block, lam=lam)

            # momentum p_n = g(r_n) v_n
            p_n = g_n @ v_n

            # (1) implicit half-step in momentum
            p = p_n.copy()
            for _ in range(int(K_p)):
                v_k = la.cho_solve(cF_n, p, check_finite=False)
                f_k = self._force_from_cache_noC(k_m_n, Dw_over_n, v_k, M_kX, D_block=D_block)
                p = p_n + 0.5 * float(eps) * f_k
            p_half = p

            # (2) implicit position update
            v_half_n = la.cho_solve(cF_n, p_half, check_finite=False)
            r = r_n + float(eps) * v_half_n
            for _ in range(int(K_q)):
                k_m_r, Dw_over_r, grad_k_r = self._rbf_cache_single(r, idx_m=idx)
                g_r, cF_r = self._metric_from_gradM(grad_k_r, M_kX, D_block=D_block, lam=lam)
                v_half_r = la.cho_solve(cF_r, p_half, check_finite=False)
                r = r_n + 0.5 * float(eps) * (v_half_n + v_half_r)
            r_np1 = r

            # (3) explicit half-step in momentum at r_{n+1}
            k_m_np1, Dw_over_np1, grad_k_np1 = self._rbf_cache_single(r_np1, idx_m=idx)
            g_np1, cF_np1 = self._metric_from_gradM(grad_k_np1, M_kX, D_block=D_block, lam=lam)
            v_mid = la.cho_solve(cF_np1, p_half, check_finite=False)
            f_np1 = self._force_from_cache_noC(k_m_np1, Dw_over_np1, v_mid, M_kX, D_block=D_block)

            p_np1 = p_half + 0.5 * float(eps) * f_np1
            v_np1 = la.cho_solve(cF_np1, p_np1, check_finite=False)

            R_next[a] = r_np1
            v_next[a] = v_np1

        if single:
            return R_next[0], v_next[0]
        return R_next, v_next

    # ============================================================
    # (NEW) Serialization + Hugging Face Hub
    # ============================================================
    def config_dict(self) -> Dict[str, Any]:
        """
        Config mínima (sin arrays grandes) para inspección/reproducibilidad.
        """
        return {
            "class": "GPLM",
            "beta": float(self.beta),
            "eps": float(self.eps),
            "sigma2": float(self.sigma2),
            "jitter": float(self.jitter),
            "seed": int(self.seed),
            "center_X": bool(self.center_X),
            "whiten_latent": bool(self.whiten_latent),
            "inducing": str(self.inducing),
            "fit_block": int(self.fit_block),
            "pred_k": None if self.pred_k is None else int(self.pred_k),
            "ann_backend": str(self.ann_backend),
            "ann_params": None if self.ann_params is None else _to_builtin(self.ann_params),
            "n_jobs": int(self.n_jobs),
            "dtype": str(self.dtype_str),
            "d_lat": int(self.d_lat),
            "D": int(self.D),
            "m": int(self.m),
        }

    def state_dict(self) -> Dict[str, Any]:
        """
        Estado completo necesario para inferencia/flow (sin datos de entrenamiento completos).
        """
        st: Dict[str, Any] = {
            "config": self.config_dict(),
            "M_mX": np.asarray(self.M_mX, dtype=np.float64),
            "Z_mx_w": np.asarray(self.Z_mx_w, dtype=np.float64),
            "mean_X": np.asarray(self.mean_X, dtype=np.float64),
            "lat_mean_x": np.asarray(self.lat_mean_x, dtype=np.float64),
            "lat_std_x": np.asarray(self.lat_std_x, dtype=np.float64),
        }
        return st

    @classmethod
    def from_state_dict(cls, st: Dict[str, Any]) -> "GPLM":
        """
        Reconstruye un objeto GPLM entrenado SIN re-entrenar.
        """
        if "config" not in st:
            raise ValueError("state_dict inválido: falta 'config'.")

        cfg = st["config"]
        # construir instancia vacía
        obj = cls.__new__(cls)

        # meta/config
        obj.beta = float(cfg["beta"])
        obj.β = obj.beta
        obj.eps = float(cfg["eps"])
        obj.ε = obj.eps
        obj.sigma2 = float(cfg["sigma2"])
        obj.σ2 = obj.sigma2
        obj.jitter = float(cfg["jitter"])
        obj.seed = int(cfg["seed"])
        obj.center_X = bool(cfg["center_X"])
        obj.whiten_latent = bool(cfg["whiten_latent"])
        obj.inducing = cfg.get("inducing", "given")
        obj.fit_block = int(cfg["fit_block"])
        obj.pred_k = cfg["pred_k"] if cfg["pred_k"] is None else int(cfg["pred_k"])
        obj.pred_κ = obj.pred_k
        obj.ann_backend = cfg.get("ann_backend", "auto")
        obj.ann_params = cfg.get("ann_params", None)
        obj.n_jobs = int(cfg.get("n_jobs", -1))

        obj.dtype_str = str(cfg.get("dtype", "float32"))
        obj.dtype = np.dtype(obj.dtype_str).type

        obj.d_lat = int(cfg["d_lat"])
        obj.D = int(cfg["D"])
        obj.m = int(cfg["m"])

        # arrays
        obj.M_mX = np.asarray(st["M_mX"], dtype=np.float64)
        obj.Z_mx_w = np.asarray(st["Z_mx_w"], dtype=np.float64)
        obj.mean_X = np.asarray(st["mean_X"], dtype=np.float64)
        obj.lat_mean_x = np.asarray(st["lat_mean_x"], dtype=np.float64)
        obj.lat_std_x = np.asarray(st["lat_std_x"], dtype=np.float64)

        # derived
        obj.Z_mx = (obj.Z_mx_w * obj.lat_std_x[None, :]) + obj.lat_mean_x[None, :]

        # ANN on inducing points (needed if pred_k is used)
        obj.ann_Z, _ = make_ann(obj.ann_backend, ann_params=obj.ann_params, n_jobs=obj.n_jobs)
        obj.ann_Z.build(obj.Z_mx_w.astype(obj.dtype, copy=False))

        # no training data kept
        obj.R_ix = None
        obj.R_iX = None
        obj.R_ix_w = None
        obj.ann_train = None
        obj.N = 0  # unknown/not needed for inference

        return obj

    def save_local(self, weights_file: str = "gplm.msgpack", config_file: str = "gplm_config.json") -> None:
        """
        Guarda:
          - weights_file: msgpack con state_dict (arrays + config)
          - config_file: JSON legible con la config
        """
        st = self.state_dict()
        cfg = st["config"]

        with open(config_file, "w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2, ensure_ascii=False)

        # msgpack (preferentemente flax.serialization)
        try:
            import flax.serialization as flax_ser  # type: ignore
        except Exception as e:
            raise RuntimeError("flax.serialization no está disponible; instale flax o use un backend alternativo.") from e

        blob = flax_ser.msgpack_serialize(st)
        with open(weights_file, "wb") as f:
            f.write(blob)

    @classmethod
    def load_local(cls, weights_file: str = "gplm.msgpack") -> "GPLM":
        """
        Carga desde msgpack (state_dict completo) y reconstruye el objeto.
        """
        try:
            import flax.serialization as flax_ser  # type: ignore
        except Exception as e:
            raise RuntimeError("flax.serialization no está disponible; instale flax o use un backend alternativo.") from e

        with open(weights_file, "rb") as f:
            blob = f.read()

        st = flax_ser.msgpack_restore(blob)
        return cls.from_state_dict(st)

    def upload_to_huggingface(
        self,
        repo_id: str,
        *,
        token: Optional[str] = None,
        weights_file: str = "gplm.msgpack",
        config_file: str = "gplm_config.json",
        repo_type: str = "model",
        revision: Optional[str] = None,
    ) -> None:
        """
        Sube (weights + config) a Hugging Face Hub.
        """
        try:
            from huggingface_hub import HfApi, create_repo  # type: ignore
        except Exception as e:
            raise RuntimeError("huggingface_hub no está instalado. Instale con `pip install huggingface_hub`.") from e

        if token is None:
            raise ValueError("token es requerido para subir al Hub (HUGGINGFACE_TOKEN/HF_TOKEN).")

        with tempfile.TemporaryDirectory() as td:
            wpath = os.path.join(td, weights_file)
            cpath = os.path.join(td, config_file)

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

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

            api.upload_file(
                path_or_fileobj=wpath,
                path_in_repo=weights_file,
                repo_id=repo_id,
                repo_type=repo_type,
                revision=revision,
            )
            api.upload_file(
                path_or_fileobj=cpath,
                path_in_repo=config_file,
                repo_id=repo_id,
                repo_type=repo_type,
                revision=revision,
            )

    @classmethod
    def download_from_huggingface(
        cls,
        repo_id: str,
        *,
        token: Optional[str] = None,
        weights_file: str = "gplm.msgpack",
        repo_type: str = "model",
        revision: Optional[str] = None,
    ) -> "GPLM":
        """
        Descarga weights_file desde el Hub y reconstruye el objeto sin re-entrenar.
        """
        try:
            from huggingface_hub import hf_hub_download  # type: ignore
        except Exception as e:
            raise RuntimeError("huggingface_hub no está instalado. Instale con `pip install huggingface_hub`.") from e

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



__all__ = ["GPLM", "InducingMode"]