"""
Data Soup — Multi-source dataset mixing for VLA training.

Implements the data soup approach from XVLA/Agibot-World:
mix multiple datasets with configurable ratios for robust training.

Supported sources:
- LeRobot datasets (HuggingFace)
- Agibot-World (bimanual manipulation)
- GR00T-Dreams synthetic data
- Custom G1 teleoperation data
- Open X-Embodiment (OXE)
- Voice commands (text instructions for language conditioning)
- Stereo4D (stereo video pairs for visual pretraining / depth)
- Cosmos DreamGen (world model synthetic data with relative actions)
"""

from __future__ import annotations

import logging
import random
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
from torch.utils.data import Dataset, ConcatDataset, WeightedRandomSampler

from neon.data.action_space import G1ActionSpace

logger = logging.getLogger(__name__)


@dataclass
class DataSourceConfig:
    """Configuration for a single data source."""

    name: str
    type: str  # "lerobot", "agibot", "dreamgen", "custom", "oxe", "voice_commands", "stereo4d", "cosmos_dreamgen"
    path: str  # HuggingFace dataset ID or local path
    weight: float = 1.0  # Sampling weight in the mixture
    max_episodes: Optional[int] = None  # Limit number of episodes
    # Action mapping
    action_key: str = "action"  # Key in the dataset for actions
    image_key: str = "observation.image"  # Key for camera images
    language_key: str = "language_instruction"  # Key for text
    proprioception_key: str = "observation.state"  # Key for joint states
    # Preprocessing
    image_size: int = 224
    fps: int = 4
    # Voice commands specific
    category_filter: Optional[List[str]] = None  # Filter by category (e.g., ["pick_place", "navigation"])
    max_samples: Optional[int] = None  # Limit total samples (for voice_commands, stereo4d)
    # Stereo4D specific
    use_stereo_pair: bool = True  # Use both left+right as multi-view input
    compute_depth: bool = False  # Compute disparity map from stereo pair
    # Cosmos DreamGen specific
    use_relative_actions: bool = True  # Use Cosmos relative action representation
    action_scaler: float = 20.0  # Scale factor for relative actions (Cosmos default)
    camera_ids: Optional[List[str]] = None  # Multi-camera IDs (e.g., ["head", "hand_0", "hand_1"])


@dataclass
class DataSoupConfig:
    """Configuration for the full data soup."""

    sources: List[DataSourceConfig] = field(default_factory=list)
    action_chunk_size: int = 16  # Number of future steps per example
    max_language_length: int = 128  # Max tokens for language
    augment_images: bool = True
    seed: int = 42


class NeonEpisode:
    """A single training episode with unified format.

    Each episode contains:
    - images: List of camera frames (PIL or numpy)
    - actions: (T, action_dim) normalized action sequence
    - language: str instruction
    - proprioception: (T, state_dim) joint state sequence
    - metadata: dict with source info
    """

    def __init__(
        self,
        images: List[Any],
        actions: np.ndarray,
        language: str = "",
        proprioception: Optional[np.ndarray] = None,
        metadata: Optional[Dict] = None,
    ):
        self.images = images
        self.actions = actions
        self.language = language
        self.proprioception = proprioception
        self.metadata = metadata or {}

    @property
    def length(self) -> int:
        return len(self.images)


class ActionMapper:
    """Maps actions from different embodiments to the G1 action space.

    Different datasets have different action representations:
    - Agibot-World: 14 DoF bimanual (7+7 arms)
    - Franka: 7 DoF + gripper
    - SO-100: 6 DoF + gripper
    - GR00T: 14 DoF bimanual
    - OXE: varies by dataset

    This mapper normalizes everything to the G1 action space.
    """

    # Known mappings from source embodiment to G1 indices
    EMBODIMENT_MAPS = {
        "agibot": {
            # Agibot has similar bimanual layout to G1 upper body
            "source_dim": 14,
            "target_indices": list(range(14)),  # Direct map to G1 arms
        },
        "franka": {
            # Franka: 7 DoF arm + gripper → map to G1 right arm
            "source_dim": 8,
            "target_indices": [7, 8, 9, 10, 11, 12, 13],  # Right arm
        },
        "so100": {
            # SO-100: 6 DoF + gripper → map to G1 right arm (approx)
            "source_dim": 7,
            "target_indices": [7, 8, 9, 10, 11, 12, 13],
        },
        "gr1": {
            # Fourier GR1: similar bimanual layout
            "source_dim": 14,
            "target_indices": list(range(14)),
        },
        "g1": {
            # G1 native: identity mapping
            "source_dim": 29,
            "target_indices": list(range(29)),
        },
    }

    def __init__(self, target_space: G1ActionSpace):
        self.target_space = target_space

    def map_actions(
        self,
        source_actions: np.ndarray,
        source_embodiment: str,
        normalize: bool = True,
    ) -> np.ndarray:
        """Map actions from source embodiment to G1 action space.

        Args:
            source_actions: (T, source_dim) actions from source
            source_embodiment: Name of source robot
            normalize: If True, normalize output to [-1, 1]

        Returns:
            (T, action_dim) actions in G1 space
        """
        mapping = self.EMBODIMENT_MAPS.get(source_embodiment.lower())
        if mapping is None:
            logger.warning(f"Unknown embodiment: {source_embodiment}, using zero-pad")
            # Zero-pad to target dim
            T = source_actions.shape[0]
            target = np.zeros((T, self.target_space.action_dim), dtype=np.float32)
            dim = min(source_actions.shape[1], self.target_space.action_dim)
            target[:, :dim] = source_actions[:, :dim]
            return target

        # Map source actions to target indices
        T = source_actions.shape[0]
        target = np.zeros((T, self.target_space.action_dim), dtype=np.float32)
        src_dim = min(source_actions.shape[1], mapping["source_dim"])
        tgt_indices = mapping["target_indices"][:src_dim]

        # Find which target indices are in the active set
        active_indices = self.target_space.active_indices
        for src_i, tgt_i in enumerate(tgt_indices):
            if tgt_i in active_indices:
                local_i = active_indices.index(tgt_i)
                if src_i < source_actions.shape[1] and local_i < target.shape[1]:
                    target[:, local_i] = source_actions[:, src_i]

        if normalize:
            target = self.target_space.normalize(target)

        return target


class DataSoupDataset(Dataset):
    """Mixed multi-source dataset for VLA training.

    Loads episodes from multiple sources, maps actions to G1 space,
    and provides (observation, action_chunk, language) tuples for training.
    """

    def __init__(self, config: DataSoupConfig, action_space: G1ActionSpace):
        self.config = config
        self.action_space = action_space
        self.action_mapper = ActionMapper(action_space)

        # Load all sources
        self.episodes: List[NeonEpisode] = []
        self.source_weights: List[float] = []
        self.source_offsets: List[Tuple[int, int]] = []  # (start, end) per source

        for source_config in config.sources:
            logger.info(f"Loading source: {source_config.name} ({source_config.type})")
            source_episodes = self._load_source(source_config)
            start = len(self.episodes)
            self.episodes.extend(source_episodes)
            end = len(self.episodes)
            self.source_offsets.append((start, end))
            self.source_weights.extend([source_config.weight] * len(source_episodes))
            logger.info(f"  Loaded {len(source_episodes)} episodes from {source_config.name}")

        # Build sample indices: each sample is (episode_idx, timestep)
        self.samples = []
        for ep_idx, ep in enumerate(self.episodes):
            for t in range(max(0, ep.length - config.action_chunk_size)):
                self.samples.append((ep_idx, t))

        logger.info(
            f"DataSoup: {len(self.episodes)} episodes, {len(self.samples)} samples "
            f"from {len(config.sources)} sources"
        )

    def _load_source(self, source_config: DataSourceConfig) -> List[NeonEpisode]:
        """Load episodes from a single source."""
        if source_config.type == "lerobot":
            return self._load_lerobot(source_config)
        elif source_config.type == "dreamgen":
            return self._load_dreamgen(source_config)
        elif source_config.type == "custom":
            return self._load_custom(source_config)
        elif source_config.type == "voice_commands":
            return self._load_voice_commands(source_config)
        elif source_config.type == "stereo4d":
            return self._load_stereo4d(source_config)
        elif source_config.type == "cosmos_dreamgen":
            return self._load_cosmos_dreamgen(source_config)
        else:
            logger.warning(f"Unknown source type: {source_config.type}")
            return []

    def _load_lerobot(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from LeRobot HuggingFace dataset format."""
        try:
            from datasets import load_dataset

            ds = load_dataset(config.path, split="train")

            # Group by episode
            episodes_data: Dict[int, List] = {}
            for item in ds:
                ep_id = item.get("episode_index", 0)
                episodes_data.setdefault(ep_id, []).append(item)

            episodes = []
            for ep_id, items in sorted(episodes_data.items()):
                if config.max_episodes and len(episodes) >= config.max_episodes:
                    break

                images = []
                actions = []
                states = []
                language = ""

                for item in items:
                    # Extract image
                    if config.image_key in item:
                        images.append(item[config.image_key])

                    # Extract action
                    if config.action_key in item:
                        actions.append(np.array(item[config.action_key], dtype=np.float32))

                    # Extract proprioception
                    if config.proprioception_key in item:
                        states.append(np.array(item[config.proprioception_key], dtype=np.float32))

                    # Extract language
                    if config.language_key in item and not language:
                        language = str(item[config.language_key])

                if not actions:
                    continue

                actions_arr = np.stack(actions)
                states_arr = np.stack(states) if states else None

                # Map actions to G1 space
                # Detect embodiment from dataset metadata
                embodiment = "franka"  # Default; could be detected from dataset
                mapped_actions = self.action_mapper.map_actions(actions_arr, embodiment)

                episodes.append(
                    NeonEpisode(
                        images=images,
                        actions=mapped_actions,
                        language=language,
                        proprioception=states_arr,
                        metadata={"source": config.name, "episode_id": ep_id},
                    )
                )

            return episodes

        except Exception as e:
            logger.error(f"Failed to load LeRobot dataset {config.path}: {e}")
            return []

    def _load_dreamgen(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from GR00T-Dreams synthetic data format.

        DreamGen output is:
        - Synthetic videos (MP4)
        - IDM-extracted actions (LeRobot format)
        """
        try:
            # DreamGen exports to LeRobot format after IDM extraction
            # So we can use the LeRobot loader
            return self._load_lerobot(config)
        except Exception as e:
            logger.error(f"Failed to load DreamGen data {config.path}: {e}")
            return []

    def _load_custom(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from custom directory structure.

        Expected format:
        path/
        ├── episode_000/
        │   ├── images/
        │   │   ├── 000.jpg
        │   │   ├── 001.jpg
        │   │   └── ...
        │   ├── actions.npy    # (T, action_dim)
        │   ├── states.npy     # (T, state_dim)
        │   └── instruction.txt
        └── episode_001/
            └── ...
        """
        from pathlib import Path

        data_dir = Path(config.path)
        if not data_dir.exists():
            logger.warning(f"Custom data dir not found: {data_dir}")
            return []

        episodes = []
        for ep_dir in sorted(data_dir.iterdir()):
            if not ep_dir.is_dir() or ep_dir.name.startswith("."):
                continue

            if config.max_episodes and len(episodes) >= config.max_episodes:
                break

            try:
                # Load images
                images_dir = ep_dir / "images"
                images = []
                if images_dir.exists():
                    from PIL import Image

                    for img_path in sorted(images_dir.glob("*.jpg")) + sorted(images_dir.glob("*.png")):
                        images.append(Image.open(img_path))

                # Load actions
                actions_path = ep_dir / "actions.npy"
                if not actions_path.exists():
                    continue
                actions = np.load(actions_path)

                # Load states
                states = None
                states_path = ep_dir / "states.npy"
                if states_path.exists():
                    states = np.load(states_path)

                # Load instruction
                language = ""
                instruction_path = ep_dir / "instruction.txt"
                if instruction_path.exists():
                    language = instruction_path.read_text().strip()

                # Map to G1 space
                mapped_actions = self.action_mapper.map_actions(actions, "g1")

                episodes.append(
                    NeonEpisode(
                        images=images,
                        actions=mapped_actions,
                        language=language,
                        proprioception=states,
                        metadata={"source": config.name, "episode_dir": str(ep_dir)},
                    )
                )
            except Exception as e:
                logger.warning(f"Failed to load episode {ep_dir}: {e}")

        return episodes

    def _load_voice_commands(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from voice commands dataset for language conditioning.

        Dataset: cagataydev/vlm-voice-commands
        Schema: {id, text, category, difficulty, voice, word_count}

        These are text-only robot instructions (50K across 10 categories).
        No images or actions — used for:
        1. Language conditioning: teach the model the distribution of natural
           language robot commands (pick/place, navigation, spatial, etc.)
        2. Audio pairing: when combined with vlm-voice-audio, provides
           text+audio pairs for the Omni backbone's native audio encoder
        3. Instruction augmentation: randomly pair voice commands with visual
           episodes from other sources to increase language diversity

        Each command becomes a 1-frame "episode" with zero actions.
        The training loop uses these for language-only loss (no action loss).
        """
        try:
            from datasets import load_dataset

            logger.info(f"Loading voice commands: {config.path}")
            ds = load_dataset(config.path, split="train", streaming=True)

            episodes = []
            count = 0

            for item in ds:
                # Category filter
                if config.category_filter:
                    if item.get("category") not in config.category_filter:
                        continue

                # Max samples limit
                if config.max_samples and count >= config.max_samples:
                    break

                text = item.get("text", "")
                if not text:
                    continue

                # Create a language-only episode (no images, zero actions)
                # These contribute to language conditioning loss only
                action_dim = self.action_space.action_dim
                episodes.append(
                    NeonEpisode(
                        images=[],  # No visual data
                        actions=np.zeros((1, action_dim), dtype=np.float32),
                        language=text,
                        proprioception=None,
                        metadata={
                            "source": config.name,
                            "type": "voice_commands",
                            "category": item.get("category", ""),
                            "difficulty": item.get("difficulty", ""),
                            "voice": item.get("voice", ""),
                            "language_only": True,  # Flag for training loop
                        },
                    )
                )
                count += 1

            logger.info(
                f"Loaded {len(episodes)} voice commands "
                f"(categories: {config.category_filter or 'all'})"
            )
            return episodes

        except Exception as e:
            logger.error(f"Failed to load voice commands {config.path}: {e}")
            return []

    def _load_stereo4d(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from stereo4D video dataset for visual pretraining.

        Dataset: cagataydev/strands-kitchen-stereo4d
        Structure: episode_XXX/stereo_left/frame_XXXXXX.png + stereo_right/...
        Loaded as imagefolder: {image: PIL} — images flattened

        Stereo4D provides:
        1. Temporal video sequences from a real kitchen environment
        2. Stereo pairs (left + right) for implicit depth understanding
        3. Scene priors: object layout, counters, spatial relations

        No actions — used for visual backbone pretraining:
        - Video temporal consistency (frame N → N+1 transitions)
        - Scene understanding (kitchen objects, workspace geometry)
        - Depth from stereo (when compute_depth=True)

        Each episode becomes a video sequence with zero actions.
        Training loop applies visual loss only (contrastive / reconstruction).

        The dataset structure on HuggingFace is:
            episode_XXX/stereo_left/frame_XXXXXX.png   (640×480)
            episode_XXX/stereo_right/frame_XXXXXX.png  (640×480)

        When loaded as imagefolder, all images are flattened.
        We reconstruct episodes from download_checksums metadata which
        contains the original paths.
        """
        try:
            from datasets import load_dataset

            logger.info(f"Loading stereo4D: {config.path}")
            ds_info = load_dataset(config.path, split="train", streaming=True)

            # The imagefolder loader flattens the directory structure.
            # We need to reconstruct episodes from the file ordering.
            # From inspection: 1038 files across ~10 episodes, ~97 frames each,
            # stereo_left + stereo_right interleaved per episode.
            #
            # Strategy: collect all images, group into chunks of consecutive
            # frames. Each chunk of left frames = one episode's video.
            # Right frames = stereo pair for that episode.

            all_images = []
            count = 0
            for item in ds_info:
                if config.max_samples and count >= config.max_samples:
                    break
                all_images.append(item["image"])
                count += 1

            if not all_images:
                logger.warning(f"No images loaded from {config.path}")
                return []

            logger.info(f"Loaded {len(all_images)} raw images from stereo4D")

            # Reconstruct episodes from the flattened image list
            # The dataset has ~10 episodes, each with stereo_left then stereo_right
            # Total 1004 images = ~10 episodes × ~50 frames × 2 cameras
            # Images are ordered: ep0/left/f0, ep0/left/f1, ..., ep0/right/f0, ...
            #
            # Each episode contributes ~97 left + ~97 right = ~194 images
            # (except shorter episodes at the end)
            #
            # Simple heuristic: split into episodes based on total count
            # and the known structure (roughly equal-sized episodes)
            total = len(all_images)
            # Each episode has left + right, with left frames first in sorted order
            # Estimate: half are left, half are right
            half = total // 2
            left_images = all_images[:half]
            right_images = all_images[half:]

            # Estimate ~10 episodes of roughly equal length
            # (from inspection: episodes 000-008 have ~97 frames, 009 has ~30)
            num_episodes = max(1, round(len(left_images) / 80))  # ~80 frames per episode avg
            frames_per_ep = len(left_images) // num_episodes

            episodes = []
            action_dim = self.action_space.action_dim

            for ep_i in range(num_episodes):
                start = ep_i * frames_per_ep
                end = start + frames_per_ep if ep_i < num_episodes - 1 else len(left_images)

                ep_left = left_images[start:end]
                if not ep_left:
                    continue

                # Subsample to target fps (dataset is dense, we want ~4fps)
                # Keep every Nth frame based on source fps vs target
                step = max(1, len(ep_left) // (config.fps * 10))  # ~10 sec episodes
                ep_frames = ep_left[::step][:32]  # Cap at 32 frames for memory

                # Optionally pair with right camera for stereo
                stereo_right = None
                if config.use_stereo_pair and ep_i * frames_per_ep < len(right_images):
                    r_start = ep_i * frames_per_ep
                    r_end = r_start + frames_per_ep if ep_i < num_episodes - 1 else len(right_images)
                    ep_right = right_images[r_start:r_end]
                    stereo_right = ep_right[::step][:32]

                T = len(ep_frames)

                episodes.append(
                    NeonEpisode(
                        images=ep_frames,
                        actions=np.zeros((T, action_dim), dtype=np.float32),  # No actions
                        language="observe kitchen scene",  # Generic scene description
                        proprioception=None,
                        metadata={
                            "source": config.name,
                            "type": "stereo4d",
                            "episode_index": ep_i,
                            "num_frames": T,
                            "has_stereo": stereo_right is not None,
                            "stereo_right": stereo_right,  # Stored in metadata for backbone
                            "visual_only": True,  # Flag: no action loss for this source
                        },
                    )
                )

            logger.info(
                f"Loaded {len(episodes)} stereo4D episodes "
                f"({sum(ep.length for ep in episodes)} total frames, "
                f"stereo_pairs={'yes' if config.use_stereo_pair else 'no'})"
            )
            return episodes

        except Exception as e:
            logger.error(f"Failed to load stereo4D {config.path}: {e}")
            return []


    def _load_cosmos_dreamgen(self, config: DataSourceConfig) -> List[NeonEpisode]:
        """Load from Cosmos-Predict2.5 DreamGen / action-conditioned data.

        Cosmos DreamGen data comes in two forms:
        1. GR1/GR00T-Dreams: HuggingFace datasets with video + text prompts
           (nvidia/GR1-100, NVIDIA/GR00T-Dreams-GR1)
        2. Bridge-style: videos/ + annotations/*.json with robot states

        Key Cosmos techniques ported here:
        - Relative end-effector actions (delta xyz/rpy in gripper frame)
        - Multi-camera conditioning (head, hand_0, hand_1)
        - Action chunking aligned to 4× temporal compression ratio
        - Scale factor of 20.0 for action deltas

        The relative action representation is crucial for cross-embodiment
        transfer. From Cosmos paper: "rel_xyz = prev_rotm.T @ (curr_xyz - prev_xyz)"
        — this makes the same delta movement mean the same thing regardless of
        workspace origin or robot geometry.

        This loader reads:
        - Videos (from HF or local path)
        - Robot state annotations (end-effector pose + gripper)
        - Text prompts (language instructions)

        And converts to NeonEpisodes using Cosmos relative actions mapped
        through our ActionMapper to the G1 action space.
        """
        try:
            import json
            from pathlib import Path

            logger.info(f"Loading Cosmos DreamGen data: {config.path}")

            # Determine if HuggingFace dataset or local path
            source_path = Path(config.path)
            episodes = []
            action_dim = self.action_space.action_dim

            if source_path.exists() and source_path.is_dir():
                # Local Cosmos-style data: videos/ + annotations/*.json
                annotations_dir = source_path / "annotations"
                videos_dir = source_path / "videos"

                if not annotations_dir.exists():
                    logger.warning(f"No annotations/ dir in {config.path}")
                    return []

                json_files = sorted(annotations_dir.glob("*.json"))
                if config.max_episodes:
                    json_files = json_files[: config.max_episodes]

                for json_path in json_files:
                    try:
                        with open(json_path) as f:
                            annotation = json.load(f)

                        # Extract relative actions using Cosmos method
                        from neon.data.relative_actions import get_action_sequence_from_states

                        rel_actions = get_action_sequence_from_states(
                            annotation,
                            fps_downsample_ratio=max(1, 30 // config.fps),  # Source typically 30fps
                            state_key=config.action_key if config.action_key != "action" else "state",
                            action_scaler=config.action_scaler,
                        )

                        # Map 7-DoF relative EE actions to G1 space
                        # Cosmos actions are [rel_xyz(3), rel_rpy(3), gripper(1)]
                        # We pad/map to full G1 action dim
                        T = len(rel_actions)
                        mapped_actions = np.zeros((T, action_dim), dtype=np.float32)
                        # Map to right arm position deltas + gripper
                        ee_dim = min(rel_actions.shape[1], action_dim)
                        mapped_actions[:, :ee_dim] = rel_actions[:, :ee_dim]

                        # Load video frames (optional — may not exist for all entries)
                        frames = []
                        video_name = json_path.stem
                        video_path = videos_dir / f"{video_name}.mp4"
                        if video_path.exists():
                            try:
                                import mediapy

                                video_array = mediapy.read_video(str(video_path))
                                step = max(1, len(video_array) // T)
                                frames = [video_array[i] for i in range(0, len(video_array), step)][:T]
                            except ImportError:
                                logger.debug("mediapy not available, skipping video frames")

                        # Text prompt
                        language = annotation.get("prompt", annotation.get("text", "robot manipulation"))

                        episodes.append(
                            NeonEpisode(
                                images=frames,
                                actions=mapped_actions,
                                language=language,
                                proprioception=None,
                                metadata={
                                    "source": config.name,
                                    "type": "cosmos_dreamgen",
                                    "action_repr": "relative_ee",
                                    "action_scaler": config.action_scaler,
                                    "original_action_dim": rel_actions.shape[1],
                                    "has_video": len(frames) > 0,
                                },
                            )
                        )
                    except Exception as e:
                        logger.debug(f"Skipping {json_path.name}: {e}")
                        continue
            else:
                # HuggingFace dataset (nvidia/GR1-100, etc.)
                from datasets import load_dataset

                logger.info(f"Loading HF Cosmos DreamGen: {config.path}")
                ds = load_dataset(config.path, split="train", streaming=True)

                count = 0
                for item in ds:
                    if config.max_episodes and count >= config.max_episodes:
                        break

                    # GR1/DreamGen datasets have varying schemas
                    # Common fields: video, text/prompt/caption
                    language = (
                        item.get("text")
                        or item.get("prompt")
                        or item.get("caption")
                        or "robot manipulation"
                    )

                    # If the dataset has state information, compute relative actions
                    if "state" in item and "continuous_gripper_state" in item:
                        from neon.data.relative_actions import get_action_sequence_from_states

                        rel_actions = get_action_sequence_from_states(
                            item, action_scaler=config.action_scaler,
                        )
                        T = len(rel_actions)
                        mapped_actions = np.zeros((T, action_dim), dtype=np.float32)
                        ee_dim = min(rel_actions.shape[1], action_dim)
                        mapped_actions[:, :ee_dim] = rel_actions[:, :ee_dim]
                    else:
                        # No state info — visual-only episode (like DreamGen bench)
                        T = 16  # Default frame count
                        mapped_actions = np.zeros((T, action_dim), dtype=np.float32)

                    episodes.append(
                        NeonEpisode(
                            images=[],  # HF streaming — frames loaded lazily
                            actions=mapped_actions,
                            language=language,
                            proprioception=None,
                            metadata={
                                "source": config.name,
                                "type": "cosmos_dreamgen",
                                "action_repr": "relative_ee" if "state" in item else "none",
                                "visual_only": "state" not in item,
                            },
                        )
                    )
                    count += 1

            logger.info(
                f"Loaded {len(episodes)} Cosmos DreamGen episodes from {config.path}"
            )
            return episodes

        except Exception as e:
            logger.error(f"Failed to load Cosmos DreamGen {config.path}: {e}")
            return []


    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """Get a single training sample.

        Returns:
            Dict with:
                - "image": Current frame (PIL Image)
                - "video_frames": Recent frames list
                - "language": Instruction text
                - "proprioception": Current joint state (action_dim,)
                - "target_actions": Future action chunk (chunk_size, action_dim)
                - "action_mask": Valid timestep mask (chunk_size,)
        """
        ep_idx, t = self.samples[idx]
        episode = self.episodes[ep_idx]
        chunk = self.config.action_chunk_size

        # Current observation
        image = episode.images[t] if t < len(episode.images) else None

        # Recent frames for video context (up to 4 past frames)
        video_start = max(0, t - 3)
        video_frames = episode.images[video_start : t + 1] if episode.images else []

        # Language instruction
        language = episode.language

        # Current proprioception
        proprio = None
        if episode.proprioception is not None and t < len(episode.proprioception):
            proprio = episode.proprioception[t].astype(np.float32)

        # Future action chunk (target)
        action_end = min(t + chunk, episode.length)
        target_actions = episode.actions[t:action_end]

        # Pad if needed
        actual_steps = target_actions.shape[0]
        if actual_steps < chunk:
            pad = np.zeros((chunk - actual_steps, target_actions.shape[1]), dtype=np.float32)
            target_actions = np.concatenate([target_actions, pad], axis=0)

        # Mask: 1 for valid steps, 0 for padding
        mask = np.zeros(chunk, dtype=np.float32)
        mask[:actual_steps] = 1.0

        return {
            "image": image,
            "video_frames": video_frames,
            "language": language,
            "proprioception": proprio,
            "target_actions": target_actions,
            "action_mask": mask,
            "metadata": episode.metadata,
            # Training flags — checked by loss computation
            "language_only": episode.metadata.get("language_only", False),
            "visual_only": episode.metadata.get("visual_only", False),
        }

    def get_sampler(self) -> WeightedRandomSampler:
        """Get a weighted sampler that respects source weights."""
        sample_weights = []
        for ep_idx, _ in self.samples:
            sample_weights.append(self.source_weights[ep_idx])

        return WeightedRandomSampler(
            weights=sample_weights,
            num_samples=len(self.samples),
            replacement=True,
        )
