from __future__ import annotations

from os import PathLike
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Tuple, cast

import torch
from torch import Tensor

from hippotorch.core.episode import Episode
from hippotorch.encoders.dual import DualEncoder
from hippotorch.encoders.protocol import MemoryEncoder
from hippotorch.memory.consolidator import Consolidator, MemoryHealthReport

if TYPE_CHECKING:  # pragma: no cover - import guard
    from hippotorch.query_api import QueryResult

from hippotorch.memory.store import MemoryStore


class HippocampalReplayBuffer:
    """Hybrid replay buffer mixing semantic, uniform, and prioritized sampling.

    When ``priority_mode`` is enabled ("proportional" or "rank"), the uniform
    portion of sampling is replaced by priority-weighted transition sampling and
    the returned batch includes an additional ``weights`` tensor for importance
    sampling correction.
    """

    def __init__(
        self,
        memory: MemoryStore,
        encoder: MemoryEncoder,
        mixture_ratio: float = 0.5,
        mixture_schedule: Optional[Callable[[int], float]] = None,
        consolidator: Optional[Consolidator] = None,
        *,
        priority_mode: str = "none",
        priority_alpha: float = 0.6,
        priority_beta: float = 0.4,
    ) -> None:
        self.memory = memory
        self.encoder = encoder
        self.mixture_ratio = mixture_ratio
        self.mixture_schedule = mixture_schedule
        self.consolidator = consolidator
        self.priority_mode = priority_mode
        self.priority_alpha = float(priority_alpha)
        self.priority_beta = float(priority_beta)
        self._total_steps = 0
        self._encoder_step = 0
        self._last_ratio = mixture_ratio
        self._sampling_stats: Dict[str, float] = {
            "semantic": 0.0,
            "uniform": 0.0,
            "retrieved_candidates": 0.0,
            "effective_ratio": mixture_ratio,
        }

    def add_episode(self, episode: Episode, encoder_step: Optional[int] = None) -> None:
        """Encode and store an episode, tracking encoder freshness."""
        episode_tensor = self._episode_to_tensor(episode).unsqueeze(0)
        # encode_key returns shape [1, embed_dim] for a single episode batch
        key = self.encoder.encode_key(episode_tensor)

        if encoder_step is None:
            self._encoder_step += 1
            step_value = self._encoder_step
        else:
            step_value = int(encoder_step)
            self._encoder_step = max(self._encoder_step, step_value)

        self.memory.write(key, [episode], encoder_step=step_value)

    def add_episodes(self, episodes: List[Episode], encoder_step: Optional[int] = None) -> None:
        """Batch-encode and store multiple episodes at once.

        Encodes all provided episodes in a single forward pass by padding them
        along the temporal dimension to the longest episode and stacking into a
        batch. The resulting key matrix is written to the underlying
        ``MemoryStore`` with a single call.
        """
        if not episodes:
            return

        # Prepare per-episode tensors (either flattened numeric or raw frames
        # depending on the encoder's ``prepare_episode`` implementation).
        ep_tensors: List[Tensor] = [self._episode_to_tensor(ep) for ep in episodes]

        # Pad along time dimension to max length and stack into [B, T, ...]
        max_len = max(t.shape[0] for t in ep_tensors)
        padded: List[Tensor] = []
        for t in ep_tensors:
            if t.shape[0] < max_len:
                pad_shape = (max_len - t.shape[0],) + tuple(t.shape[1:])
                pad = torch.zeros(*pad_shape, device=t.device, dtype=t.dtype)
                t = torch.cat([t, pad], dim=0)
            padded.append(t)
        batch = torch.stack(padded, dim=0)

        # Encode in a single forward pass
        keys = self.encoder.encode_key(batch)

        # Update encoder step bookkeeping
        if encoder_step is None:
            # Match semantics of repeated single insertions: advance by count
            self._encoder_step += len(episodes)
            step_value = self._encoder_step
        else:
            step_value = int(encoder_step)
            self._encoder_step = max(self._encoder_step, step_value)

        # Single write into memory
        self.memory.write(keys, episodes, encoder_step=step_value)

    @property
    def size(self) -> int:
        """Number of stored episodes."""

        return len(self.memory)

    def attach_consolidator(self, consolidator: Consolidator) -> None:
        """Attach a consolidator post-construction."""

        self.consolidator = consolidator

    def sample(
        self, batch_size: int, query_state: Optional[Tensor] = None, top_k: int = 5
    ) -> Dict[str, Tensor]:
        if len(self.memory) == 0:
            raise ValueError("Memory is empty; add episodes before sampling.")

        ratio = self._current_ratio()
        self._last_ratio = ratio
        n_semantic = min(int(batch_size * ratio), batch_size)
        n_uniform = batch_size - n_semantic

        samples: List[Tuple[Episode, int]] = []
        is_weights: List[Tensor] = []
        retrieved_candidates = 0

        if n_semantic > 0 and query_state is not None:
            query = self._prepare_query(query_state)
            _, retrieved, _ = self.memory.retrieve(
                query,
                top_k=top_k,
                encoder=self.encoder,
                current_step=self._encoder_step,
                lazy_refresh=True,
            )
            retrieved_eps = retrieved[0] if retrieved else []
            retrieved_candidates = len(retrieved_eps)
            if retrieved_eps:
                samples.extend(self._scatter_sample(retrieved_eps, n_semantic))
            else:
                n_uniform += n_semantic

        if n_uniform > 0:
            if self.priority_mode in ("proportional", "rank"):
                pr_samples, weights, _ = self.memory.sample_priority_transitions(
                    n_uniform,
                    mode=self.priority_mode,
                    alpha=self.priority_alpha,
                    beta=self.priority_beta,
                )
                samples.extend(pr_samples)
                if weights.numel() > 0:
                    is_weights.append(weights)
            else:
                uniform_eps = self.memory.sample_uniform(n_uniform)
                for episode in uniform_eps:
                    if len(episode) <= 1:
                        continue
                    idx_val = int(torch.randint(0, len(episode) - 1, (1,)).item())
                    samples.append((episode, idx_val))

        if not samples:
            raise ValueError("Unable to draw samples from memory.")

        self._total_steps += len(samples)
        total_drawn = len(samples)
        effective_ratio = (total_drawn - n_uniform) / total_drawn if total_drawn else 0.0
        self._sampling_stats = {
            "semantic": float(total_drawn - n_uniform),
            "uniform": float(n_uniform),
            "retrieved_candidates": float(retrieved_candidates),
            "effective_ratio": float(effective_ratio),
        }
        batch = self._samples_to_batch(samples[:batch_size])
        if self.priority_mode in ("proportional", "rank"):
            # If some samples came from prioritized sampling, provide IS weights.
            if is_weights:
                w = torch.cat(is_weights, dim=0)
                batch["weights"] = w[: batch["states"].shape[0]].to(batch["states"].device)
            else:
                batch["weights"] = torch.ones(
                    batch["states"].shape[0], device=batch["states"].device
                )
        return batch

    def sample_with_intrinsic(
        self,
        batch_size: int,
        *,
        query_state: Optional[Tensor] = None,
        top_k: int = 5,
        intrinsic_top_k: int = 5,
        intrinsic_scale: float = 1.0,
    ) -> Dict[str, Tensor]:
        """Sample a batch and attach an intrinsic curiosity reward.

        Computes inverse-distance novelty for each sampled state against memory keys.
        """
        batch = self.sample(batch_size=batch_size, query_state=query_state, top_k=top_k)
        try:
            from .intrinsic import IntrinsicConsolidator

            ic = IntrinsicConsolidator(
                memory=self.memory,
                encoder=self.encoder,
                target_similarity=1.0,
                scale=float(intrinsic_scale),
                top_k=int(intrinsic_top_k),
            )
            intrinsic = []
            for s in batch["states"]:
                out = ic.compute(torch.cat([s, torch.zeros(2, device=s.device)]))
                intrinsic.append(out.intrinsic_reward)
            batch["intrinsic"] = torch.stack(intrinsic)
        except Exception:
            pass
        return batch

    def sample_windows(
        self,
        batch_size: int,
        *,
        window_size: int = 10,
        query_state: Optional[Tensor] = None,
        top_k: int = 5,
    ) -> Dict[str, Tensor]:
        """Draw a batch of temporal windows for POMDP agents.

        Returns a dictionary with padded tensors: states, actions, rewards, dones, lengths.
        """
        if len(self.memory) == 0:
            raise ValueError("Memory is empty; add episodes before sampling.")

        ratio = self._current_ratio()
        n_semantic = min(int(batch_size * ratio), batch_size)
        n_uniform = batch_size - n_semantic

        windows: List[Episode] = []

        if n_semantic > 0 and query_state is not None:
            query = self._prepare_query(query_state)
            _, retrieved, _ = self.memory.retrieve(
                query,
                top_k=top_k,
                encoder=self.encoder,
                current_step=self._encoder_step,
                lazy_refresh=True,
            )
            retrieved_eps = retrieved[0] if retrieved else []
            if retrieved_eps:
                windows.extend(
                    self._scatter_sample_windows(retrieved_eps, n_semantic, window_size=window_size)
                )
            else:
                n_uniform += n_semantic

        if n_uniform > 0:
            uniform_eps = self.memory.sample_uniform(n_uniform)
            windows.extend(
                self._scatter_sample_windows(uniform_eps, n_uniform, window_size=window_size)
            )

        if not windows:
            raise ValueError("Unable to draw sequence windows from memory.")

        return self._windows_to_batch(windows[:batch_size])

    def consolidate(
        self,
        *,
        steps: int = 1,
        batch_size: int = 32,
        refresh_keys: bool = True,
        report_quality: bool = False,
        tb_writer=None,
        tb_log_every: int = 0,
    ) -> Dict[str, float]:
        """Run consolidation through the attached Consolidator."""

        if self.consolidator is None:
            raise RuntimeError("No consolidator attached to replay buffer")

        metrics = self.consolidator.sleep(
            self.memory,
            steps=steps,
            batch_size=batch_size,
            refresh_keys=refresh_keys,
            report_quality=report_quality,
            tb_writer=tb_writer,
            tb_log_every=tb_log_every,
        )
        return metrics

    # ------------------------------------------------------------------
    def save(self, path: PathLike[str] | str) -> str:
        """Save buffer snapshot (memory + encoder) to ``path`` for sharing."""

        target = Path(path)
        target.parent.mkdir(parents=True, exist_ok=True)
        snapshot = self._to_snapshot()
        torch.save(snapshot, target)
        return str(target)

    @classmethod
    def load(
        cls, path: PathLike[str] | str, *, device: str | torch.device = "cpu"
    ) -> "HippocampalReplayBuffer":
        """Load a buffer snapshot saved with :meth:`save`."""

        payload = cast(Dict[str, object], torch.load(Path(path), map_location=device))
        return cls._from_snapshot(payload, device=device)

    def _to_snapshot(self) -> Dict[str, object]:
        return {
            "snapshot_version": 1,
            "buffer": {
                "mixture_ratio": self.mixture_ratio,
                "priority_mode": self.priority_mode,
                "priority_alpha": float(self.priority_alpha),
                "priority_beta": float(self.priority_beta),
            },
            "memory": self.memory.to_checkpoint(),
            "encoder": self._serialize_encoder(),
        }

    @classmethod
    def _from_snapshot(
        cls, payload: Mapping[str, object], *, device: str | torch.device
    ) -> "HippocampalReplayBuffer":
        memory = MemoryStore.from_checkpoint(
            cast(Mapping[str, object], payload["memory"]),
            device=device,
        )
        encoder = cls._deserialize_encoder(
            cast(Mapping[str, object], payload["encoder"]),
            device=device,
        )
        buffer_cfg = cast(Dict[str, Any], payload.get("buffer", {}))
        mixture = float(cast(Any, buffer_cfg.get("mixture_ratio", 0.0)))
        priority_mode = cast(str, buffer_cfg.get("priority_mode", "none"))
        priority_alpha = float(cast(Any, buffer_cfg.get("priority_alpha", 0.6)))
        priority_beta = float(cast(Any, buffer_cfg.get("priority_beta", 0.4)))
        return cls(
            memory=memory,
            encoder=encoder,
            mixture_ratio=mixture,
            priority_mode=priority_mode,
            priority_alpha=priority_alpha,
            priority_beta=priority_beta,
        )

    def _serialize_encoder(self) -> Dict[str, object]:
        def _strip_module_prefix(sd: Mapping[str, Any]) -> Dict[str, Any]:
            """Return a copy of state dict without any '.module.' prefixes.

            Handles common cases where submodules like 'online'/'target' are wrapped
            in DataParallel, which introduces 'module.' into parameter names.
            """
            new_sd: Dict[str, Any] = {}
            for k, v in sd.items():
                nk = k
                # Remove only the first occurrence of '.module.' after a prefix
                nk = nk.replace(".module.", ".")
                # Also handle top-level 'module.' prefix conservatively
                if nk.startswith("module."):
                    nk = nk[len("module.") :]
                new_sd[nk] = v
            return new_sd

        encoder = self.encoder
        if isinstance(encoder, DualEncoder):
            # Narrow types for mypy
            from torch import nn as _nn

            input_proj = encoder.online.input_proj
            input_dim = int(getattr(input_proj, "in_features", 0))
            first_linear = cast(_nn.Linear, encoder.projector[0])
            embed_dim = int(first_linear.out_features)
            return {
                "type": "DualEncoder",
                "config": {
                    "input_dim": input_dim,
                    "embed_dim": embed_dim,
                    "momentum": float(encoder.momentum),
                    "refresh_interval": int(encoder.refresh_interval),
                },
                # Save with DP prefixes stripped for portability
                "state_dict": _strip_module_prefix(encoder.state_dict()),
            }
        raise TypeError("Buffer serialization currently supports DualEncoder instances only.")

    @staticmethod
    def _deserialize_encoder(
        payload: Mapping[str, object], *, device: str | torch.device
    ) -> DualEncoder:
        enc_type = payload.get("type")
        if enc_type != "DualEncoder":
            raise ValueError(f"Unsupported encoder type {enc_type!r} in snapshot.")
        config = cast(Dict[str, Any], payload.get("config", {}))
        encoder = DualEncoder(
            input_dim=int(cast(Any, config.get("input_dim", 0))),
            embed_dim=int(cast(Any, config.get("embed_dim", 128))),
            momentum=float(cast(Any, config.get("momentum", 0.995))),
            refresh_interval=int(cast(Any, config.get("refresh_interval", 10000))),
        )
        state_dict = cast(Mapping[str, Any], payload.get("state_dict", {}))

        # Accept state dicts saved from DataParallel (with 'module.' prefixes)
        def _strip_module_prefix(sd: Mapping[str, Any]) -> Dict[str, Any]:
            new_sd: Dict[str, Any] = {}
            for k, v in sd.items():
                nk = k.replace(".module.", ".")
                if nk.startswith("module."):
                    nk = nk[len("module.") :]
                new_sd[nk] = v
            return new_sd

        try:
            encoder.load_state_dict(state_dict)
        except Exception:
            # Retry after stripping DP prefixes
            cleaned = _strip_module_prefix(state_dict)
            encoder.load_state_dict(cleaned)
        encoder.to(device)
        return encoder

    def health_report(self, sample_size: int = 64, top_k: int = 5) -> MemoryHealthReport:
        """Return a :class:`MemoryHealthReport` describing memory health."""

        return MemoryHealthReport.from_memory(
            self.memory,
            self.encoder,
            sample_size=sample_size,
            top_k=top_k,
            step=self._encoder_step,
        )

    def query_batch(self, query_states: Tensor, *, top_k: int = 5) -> "QueryResult":
        """Batch retrieval helper returning a :class:`QueryResult`."""

        from hippotorch.query_api import query_batch as _query_batch

        return _query_batch(
            query_states,
            buffer=self,
            top_k=top_k,
            return_scores=True,
            return_attention=False,
        )

    def save_to_hub(
        self,
        repo_id: str,
        *,
        token: Optional[str] = None,
        revision: str = "main",
        private: bool = False,
        exist_ok: bool = True,
        filename: str = "memory.pt",
        workdir: Optional[str] = None,
        readme: Optional[str] = None,
        backend: str | None = None,
        local_dir: Optional[str] = None,
    ) -> PathLike[str] | str:
        """Upload the attached memory to the HuggingFace Hub.

        This is a thin convenience wrapper around
        :func:`hippotorch.utils.hub.push_memory_to_hub` that keeps the buffer API
        ergonomic for transfer learning workflows.
        """

        from hippotorch.utils.hub import push_memory_to_hub

        return push_memory_to_hub(
            self,
            repo_id,
            token=token,
            revision=revision,
            private=private,
            exist_ok=exist_ok,
            filename=filename,
            local_dir=local_dir or workdir,
            readme=readme,
            backend=backend,
        )

    def load_memory_from_hub(
        self,
        repo_id: str,
        *,
        filename: str = "memory.pt",
        revision: str = "main",
        token: Optional[str] = None,
        device: Optional[str | torch.device] = None,
        cache_dir: Optional[str] = None,
        backend: str | None = None,
        local_dir: Optional[str] = None,
    ) -> MemoryStore:
        """Replace the underlying memory with a checkpoint from the Hub."""

        from hippotorch.utils.hub import load_memory_from_hub

        target_device = device or getattr(self.memory, "device", torch.device("cpu"))
        loaded = load_memory_from_hub(
            repo_id,
            filename=filename,
            revision=revision,
            token=token,
            device=target_device,
            cache_dir=cache_dir,
            backend=backend,
            local_dir=local_dir,
        )
        mem = cast(MemoryStore, loaded)
        self._validate_memory_compatibility(mem)
        self.memory = mem
        return mem

    def _validate_memory_compatibility(self, memory: MemoryStore) -> None:
        """Ensure loaded memory matches encoder and buffer expectations."""

        # Best-effort detection of encoder output dimension for compatibility checks
        encoder_embed = None
        try:
            from torch import nn as _nn

            proj = getattr(self.encoder, "projector", None)
            if hasattr(proj, "out_features"):
                encoder_embed = int(getattr(proj, "out_features"))
            elif isinstance(proj, _nn.Sequential) and len(proj) > 0:
                first = proj[0]
                if isinstance(first, _nn.Linear):
                    encoder_embed = int(first.out_features)
        except Exception:
            encoder_embed = None
        if encoder_embed is not None and memory.embed_dim != int(encoder_embed):
            raise ValueError(
                (
                    "Loaded memory embed_dim "
                    f"{memory.embed_dim} does not match encoder embed_dim {encoder_embed}."
                )
            )

        if hasattr(self.memory, "embed_dim") and memory.embed_dim != self.memory.embed_dim:
            raise ValueError(
                (
                    "Loaded memory embed_dim "
                    f"{memory.embed_dim} does not match buffer embed_dim {self.memory.embed_dim}."
                )
            )

    def _current_ratio(self) -> float:
        if self.mixture_schedule is None:
            return self.mixture_ratio
        return float(self.mixture_schedule(self._total_steps))

    def _prepare_query(self, state: Tensor) -> Tensor:
        prepare_fn = getattr(self.encoder, "prepare_query", None)
        if callable(prepare_fn):
            prepared = prepare_fn(state)
            return self.encoder.encode_query(prepared)

        if state.dim() == 1:
            state = state.unsqueeze(0).unsqueeze(0)
        elif state.dim() == 2:
            state = state.unsqueeze(0)
        # Align feature dimension with encoder input expectations (obs + action + reward)
        try:
            in_feats = self.encoder.online.input_proj.in_features  # type: ignore[attr-defined]
        except Exception:
            in_feats = state.shape[-1]
        d = state.shape[-1]
        if d < in_feats:
            pad = torch.zeros(
                *state.shape[:-1], in_feats - d, device=state.device, dtype=state.dtype
            )
            state = torch.cat([state, pad], dim=-1)
        elif d > in_feats:
            state = state[..., :in_feats]
        return self.encoder.encode_query(state)

    def _scatter_sample(self, episodes: List[Episode], n_samples: int) -> List[Tuple[Episode, int]]:
        samples: List[Tuple[Episode, int]] = []
        per_episode = max(1, n_samples // len(episodes))
        for episode in episodes:
            if len(episode) <= 1:
                continue
            indices = torch.randint(0, len(episode) - 1, (per_episode,))
            for idx in indices.tolist():
                samples.append((episode, idx))
                if len(samples) >= n_samples:
                    return samples
        return samples

    def _episode_to_tensor(self, episode: Episode) -> Tensor:
        prepare_fn = getattr(self.encoder, "prepare_episode", None)
        if callable(prepare_fn):
            tensor = prepare_fn(episode)
            if not isinstance(tensor, torch.Tensor):
                raise TypeError("prepare_episode must return a torch.Tensor")
            return tensor

        states = episode.states
        if states.dim() > 2:
            states = states.reshape(states.shape[0], -1)
        if states.dtype != episode.actions.dtype:
            states = states.to(episode.actions.dtype)
        rewards = episode.rewards.unsqueeze(-1).to(episode.actions.dtype)
        return torch.cat([states, episode.actions, rewards], dim=-1)

    def _samples_to_batch(self, samples: List[Tuple[Episode, int]]) -> Dict[str, Tensor]:
        states, actions, rewards, next_states, dones = [], [], [], [], []
        for episode, idx in samples:
            states.append(episode.states[idx])
            actions.append(episode.actions[idx])
            rewards.append(episode.rewards[idx])
            done_value = (
                episode.dones[idx]
                if episode.dones is not None
                else torch.zeros((), dtype=torch.bool, device=episode.states.device)
            )
            dones.append(done_value)
            if idx + 1 < len(episode):
                next_states.append(episode.states[idx + 1])
            else:
                next_states.append(episode.states[idx])

        return {
            "states": torch.stack(states),
            "actions": torch.stack(actions),
            "rewards": torch.stack(rewards),
            "next_states": torch.stack(next_states),
            "dones": torch.stack(dones),
        }

    # --- Hub portability helpers ---
    # (hub helper methods defined above; duplicates removed)

    def _scatter_sample_windows(
        self, episodes: List[Episode], n_samples: int, *, window_size: int
    ) -> List[Episode]:
        """Sample random windows of length `window_size` from episodes."""
        if not episodes:
            return []
        out: List[Episode] = []
        attempts = 0
        while len(out) < n_samples and attempts < n_samples * 5:
            idx = int(torch.randint(0, len(episodes), (1,)).item())
            ep = episodes[idx]
            if len(ep) <= window_size:
                attempts += 1
                continue
            start = int(torch.randint(0, len(ep) - window_size, (1,)).item())
            out.append(ep.slice(start, start + window_size))
            attempts += 1
        return out

    def _windows_to_batch(self, windows: List[Episode]) -> Dict[str, Tensor]:
        """Pad variable-length windows and stack into batch tensors.

        Returns a dict with keys: states [B, T, Ds], actions [B, T, Da],
        rewards [B, T], dones [B, T], lengths [B].
        """
        if not windows:
            raise ValueError("No windows provided")
        lengths = torch.tensor([len(w) for w in windows], dtype=torch.long)
        max_len = int(lengths.max().item())
        states_b, actions_b, rewards_b, dones_b = [], [], [], []
        for w in windows:
            s, a, r = w.states, w.actions, w.rewards
            d = (
                w.dones
                if w.dones is not None
                else torch.zeros(len(w), dtype=torch.bool, device=s.device)
            )
            if len(w) < max_len:
                s = torch.cat(
                    [
                        s,
                        torch.zeros(
                            max_len - len(w),
                            *s.shape[1:],
                            device=s.device,
                            dtype=s.dtype,
                        ),
                    ],
                    dim=0,
                )
                a = torch.cat(
                    [
                        a,
                        torch.zeros(
                            max_len - len(w),
                            *a.shape[1:],
                            device=a.device,
                            dtype=a.dtype,
                        ),
                    ],
                    dim=0,
                )
                r = torch.cat(
                    [r, torch.zeros(max_len - len(w), device=r.device, dtype=r.dtype)],
                    dim=0,
                )
                d = torch.cat(
                    [d, torch.zeros(max_len - len(w), device=d.device, dtype=d.dtype)],
                    dim=0,
                )
            states_b.append(s)
            actions_b.append(a)
            rewards_b.append(r)
            dones_b.append(d)
        return {
            "states": torch.stack(states_b),
            "actions": torch.stack(actions_b),
            "rewards": torch.stack(rewards_b),
            "dones": torch.stack(dones_b),
            "lengths": lengths,
        }

    def stats(self) -> Dict[str, float]:
        """Return diagnostic statistics about sampling and storage."""

        stats = {
            "memory_size": float(len(self.memory)),
            "semantic_fraction": float(self._last_ratio),
            "encoder_step": float(self._encoder_step),
            "sampled_steps": float(self._total_steps),
        }
        staleness = self.memory.staleness_stats(current_step=self._encoder_step)
        stats.update({f"staleness_{k}": v for k, v in staleness.items()})
        stats.update(self._sampling_stats)
        return stats

    # --- Prioritization API ---
    def update_priorities(
        self,
        indices: Tensor | List[Tuple[int, int]],
        td_errors: Tensor,
        *,
        eps: float = 1e-6,
    ) -> None:
        """Update per-transition priorities in the underlying memory store.

        Args:
            indices: Sequence or tensor of (episode_idx, step_idx) pairs.
            td_errors: Absolute TD-error values per transition.
            eps: Small constant to keep priorities positive.
        """
        self.memory.update_priorities(indices, td_errors, eps=eps)

    def sampling_stats(self) -> Dict[str, float]:
        """Expose the most recent sampling diagnostics."""

        return dict(self._sampling_stats)


__all__ = ["HippocampalReplayBuffer"]
