"""Offline goal-conditioned samples drawn from the cached PushT latents.

A sample is a 3-frame latent history, the two action blocks that connect
those frames, a goal latent taken from 1-``max_offset`` transitions ahead in
the *same* episode, and the real action block leaving the current frame (used
only to fit the behavior-density model, never as a controller target).

Everything is in latent space, so this never touches the 46 GB image file.
"""

import json
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset

NUM_CONTEXT = 3
FRAMESKIP = 5


class LatentGoalDataset(Dataset):
    """Latent history/goal pairs sampled from cached expert trajectories.

    Args:
        root: Directory written by ``scripts/encode_latents.py``.
        max_offset: Largest goal distance, in world-model transitions. Set by
            the training curriculum (1-2, then 1-3, then 1-5).
        episodes: Episode indices this split may draw from.
        horizon: Plan horizon, reserved so a full plan stays inside the episode.
    """

    def __init__(
        self, root, max_offset=5, episodes=None, horizon=5, in_memory=True
    ):
        root = Path(root)
        stats = json.loads((root / 'stats.json').read_text())
        self.latent_dim = stats['latent_dim']

        # 2.34M x 192 fp16 is ~0.9 GB, so keeping it resident beats paging a
        # memmap with random access from every dataloader worker.
        mode = None if in_memory else 'r'
        self.latents = np.load(root / 'latents.npy', mmap_mode=mode)
        self.actions = np.load(root / 'actions.npy', mmap_mode=mode)
        self.lengths = np.load(root / 'ep_len.npy')
        self.offsets = np.load(root / 'ep_offset.npy')

        self.action_mean = np.asarray(stats['action_mean'], dtype=np.float32)
        self.action_std = np.asarray(stats['action_std'], dtype=np.float32)

        self.horizon = horizon
        self.set_max_offset(max_offset)

        if episodes is None:
            episodes = np.arange(len(self.lengths))
        self.episodes = np.asarray(episodes)

        # Flat (episode, start) index as arrays — a Python list of 2M tuples
        # gets re-pickled into every worker and dominates startup cost.
        # A clip needs its N context frames *and* at least one further
        # transition to place a goal in; without that reserve the last clip of
        # each episode silently takes its goal from the next episode.
        counts = (self.lengths[self.episodes] - NUM_CONTEXT * FRAMESKIP).clip(
            min=0
        )
        self.clip_episode = np.repeat(self.episodes, counts)
        self.clip_start = np.concatenate(
            [np.arange(c) for c in counts if c > 0]
        ).astype(np.int64)

    def set_max_offset(self, max_offset):
        """Update the curriculum's furthest goal, in world-model transitions."""
        self.max_offset = max_offset

    def __len__(self):
        return len(self.clip_start)

    def _block(self, base, local_step):
        """Normalized action block leaving the frame at ``local_step``.

        Matches ``swm.data.Dataset.__getitem__``: pixels are strided by
        frameskip while actions are not, so the block leaving frame ``k`` is
        the raw actions ``[start + 5k, start + 5k + 5)``.
        """
        lo = base + local_step
        raw = np.asarray(self.actions[lo : lo + FRAMESKIP], dtype=np.float32)
        return ((raw - self.action_mean) / self.action_std).reshape(-1)

    def __getitem__(self, idx):
        ep = int(self.clip_episode[idx])
        start = int(self.clip_start[idx])
        base = int(self.offsets[ep])
        length = int(self.lengths[ep])

        # context frames, strided by frameskip
        rows = [base + start + i * FRAMESKIP for i in range(NUM_CONTEXT)]
        ctx = np.asarray(self.latents[rows], dtype=np.float32)

        # blocks between the context frames: block i leaves context frame i
        past = np.stack(
            [
                self._block(base, start + i * FRAMESKIP)
                for i in range(NUM_CONTEXT - 1)
            ]
        )

        # goal: 1..max_offset transitions past the current frame, clamped to
        # what is left in the episode
        current = start + (NUM_CONTEXT - 1) * FRAMESKIP
        reach = (length - 1 - current) // FRAMESKIP
        high = min(self.max_offset, reach)
        offset = int(torch.randint(1, high + 1, (1,)).item())
        goal = np.asarray(
            self.latents[base + current + offset * FRAMESKIP], dtype=np.float32
        )

        # the real block leaving the current frame — density model only
        real = self._block(base, current)

        return {
            'context': torch.from_numpy(ctx),
            'past_actions': torch.from_numpy(past),
            'goal': torch.from_numpy(goal),
            'real_action': torch.from_numpy(real),
            'goal_offset': torch.tensor(offset, dtype=torch.long),
        }


def split_episodes(num_episodes, val_fraction=0.05, seed=0):
    """Deterministic train/val split over episode indices."""
    rng = np.random.default_rng(seed)
    perm = rng.permutation(num_episodes)
    n_val = max(1, int(num_episodes * val_fraction))
    return perm[n_val:], perm[:n_val]