"""
LatentDataset: Load precomputed ctx and target latents for Context-as-Memory dataset.

Use after running precompute_ctx_target_latents.py. Returns samples with:
- context_latents: (1, C, K, H//8, W//8) - 1 latent per context frame
- target_latents: (1, C, T, H//8, W//8) - 1 latent per 4 target frames
- prompt, video_name, start_frame, end_frame, actions

Compatible with training that uses precomputed latents instead of encoding on the fly.
"""

import json
import os
import warnings

import torch


class LatentDataset(torch.utils.data.Dataset):
    """
    Dataset that loads precomputed ctx and target latents.
    """

    def __init__(
        self,
        latent_dir,
        metadata_path=None,
        action_base_path=None,
        repeat=1,
        num_frames=81,
        context_frames=5,
        target_frames_per_latent=4,
    ):
        """
        Args:
            latent_dir: Directory containing ctx_latents/ and target_latents/ subdirs.
            metadata_path: Optional. If provided, used to get total_samples and validate.
            action_base_path: Base path for action JSON files (for loading actions if not in .pt).
            repeat: Dataset repeat factor.
            num_frames: Expected num_frames per segment.
            context_frames: Number of context frames (K).
            target_frames_per_latent: Target: 1 latent per N frames.
        """
        self.latent_dir = latent_dir
        self.ctx_dir = os.path.join(latent_dir, "ctx_latents")
        self.target_dir = os.path.join(latent_dir, "target_latents")
        self.action_base_path = action_base_path or latent_dir
        self.repeat = repeat
        self.num_frames = num_frames
        self.context_frames = context_frames
        self.target_frames_per_latent = target_frames_per_latent

        # Infer valid indices from existing files (both ctx and target must exist)
        self._indices = []
        if os.path.isdir(self.ctx_dir) and os.path.isdir(self.target_dir):
            ctx_files = {f.replace(".pt", "") for f in os.listdir(self.ctx_dir) if f.endswith(".pt")}
            target_files = {f.replace(".pt", "") for f in os.listdir(self.target_dir) if f.endswith(".pt")}
            common = sorted([int(x) for x in ctx_files & target_files])
            self._indices = common
        if not self._indices:
            meta_path = os.path.join(latent_dir, "metadata_precompute.json")
            if os.path.isfile(meta_path):
                with open(meta_path) as f:
                    meta = json.load(f)
                self._total = meta.get("total_samples", 0)
                self._indices = list(range(self._total))
            else:
                self._total = 0
        else:
            self._total = len(self._indices)

    def __len__(self):
        return self._total * self.repeat

    def __getitem__(self, idx):
        real_idx = idx % self._total
        if self._indices is not None:
            real_idx = self._indices[real_idx]

        ctx_path = os.path.join(self.ctx_dir, f"{real_idx:08d}.pt")
        target_path = os.path.join(self.target_dir, f"{real_idx:08d}.pt")

        if not os.path.isfile(ctx_path) or not os.path.isfile(target_path):
            warnings.warn(f"Latent files not found for idx {real_idx}. Returning None.")
            return None

        ctx_data = torch.load(ctx_path, map_location="cpu", weights_only=True)
        target_data = torch.load(target_path, map_location="cpu", weights_only=True)

        ctx_latent = ctx_data["latent"]
        target_latent = target_data["latent"]

        # Ensure batch dimension: (C, K, H, W) -> (1, C, K, H, W)
        if ctx_latent.dim() == 4:
            ctx_latent = ctx_latent.unsqueeze(0)
        if target_latent.dim() == 4:
            target_latent = target_latent.unsqueeze(0)

        out = {
            "context_latents": ctx_latent,
            "target_latents": target_latent,
            "prompt": ctx_data.get("prompt", ""),
            "video_name": ctx_data.get("video_name"),
            "start_frame": ctx_data.get("start_frame"),
            "end_frame": ctx_data.get("end_frame"),
        }
        if "actions" in ctx_data and ctx_data["actions"] is not None:
            out["actions"] = ctx_data["actions"]
        elif "actions" in target_data and target_data["actions"] is not None:
            out["actions"] = target_data["actions"]

        return out


def get_latent_dataset_args(latent_dir, action_base_path=None, **kwargs):
    """Build argparse.Namespace for LatentDataset from precompute metadata."""
    meta_path = os.path.join(latent_dir, "metadata_precompute.json")
    if not os.path.isfile(meta_path):
        return None
    with open(meta_path) as f:
        meta = json.load(f)
    from argparse import Namespace
    return Namespace(
        latent_dir=latent_dir,
        action_base_path=action_base_path or meta.get("dataset_base_path", latent_dir),
        num_frames=meta.get("num_frames", 81),
        context_frames=meta.get("context_frames", 5),
        target_frames_per_latent=meta.get("target_frames_per_latent", 4),
        **kwargs,
    )