"""Lazy prompt-sharded access for large Predictor training datasets."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open
from torch.utils.data import Dataset


class LazyLayer17Dataset(Dataset):
    """Read only the tensors needed by one Layer-17 sample on demand."""

    def __init__(
        self,
        root: str | Path,
        layer_id: int = 17,
        include_previous_hidden: bool = True,
    ) -> None:
        self.root = Path(root).resolve()
        self.layer_id = int(layer_id)
        self.include_previous_hidden = bool(include_previous_hidden)

    def __len__(self) -> int:
        return 1000

    def __getitem__(self, index: tuple[int, int, int]) -> dict[str, Any]:
        prompt_id, chunk, target_step = (int(value) for value in index)
        if not 0 <= prompt_id < 1000:
            raise ValueError(index)
        if chunk < 1 or chunk > 6 or target_step < 1 or target_step > 3:
            raise ValueError(index)

        anchor_step = target_step - 1
        prompt_dir = self.root / f"prompt_{prompt_id:04d}"
        trajectory_path = prompt_dir / "trajectory.safetensors"
        context_path = prompt_dir / "chunk0_context" / "trajectory.safetensors"

        with safe_open(trajectory_path, framework="pt", device="cpu") as handle:
            prefix = f"chunk_{chunk:02d}_step_{target_step:02d}"
            anchor_prefix = f"chunk_{chunk:02d}_step_{anchor_step:02d}"
            noisy = handle.get_tensor(f"{prefix}_noisy_latent").squeeze(0)
            timestep = handle.get_tensor(f"{prefix}_timestep").squeeze(0)
            anchor_timestep = handle.get_tensor(
                f"{anchor_prefix}_timestep"
            ).squeeze(0)
            anchor_hidden = handle.get_tensor(
                f"{anchor_prefix}_final_hidden"
            ).squeeze(0)
            target_hidden = handle.get_tensor(
                f"{prefix}_final_hidden"
            ).squeeze(0)
            target_flow = handle.get_tensor(f"{prefix}_flow").squeeze(0)
            if self.include_previous_hidden and chunk > 1:
                previous_hidden = handle.get_tensor(
                    f"chunk_{chunk - 1:02d}_step_{target_step:02d}_final_hidden"
                ).squeeze(0)
            else:
                previous_hidden = None
        if self.include_previous_hidden and previous_hidden is None:
            with safe_open(context_path, framework="pt", device="cpu") as handle:
                previous_hidden = handle.get_tensor(
                    f"chunk_00_step_{target_step:02d}_final_hidden"
                ).squeeze(0)
        elif not self.include_previous_hidden:
            # Keep the common batch interface while the DisCa model physically
            # omits this channel from its fusion layer.
            previous_hidden = torch.zeros_like(anchor_hidden)

        prefeature_path = (
            prompt_dir
            / "clean_prefeatures"
            / f"block_{self.layer_id:02d}.safetensors"
        )
        context_prefeature_path = (
            prompt_dir
            / "chunk0_context"
            / "clean_prefeatures"
            / f"block_{self.layer_id:02d}.safetensors"
        )
        history_parts = []
        with safe_open(
            context_prefeature_path, framework="pt", device="cpu"
        ) as context_handle, safe_open(
            prefeature_path, framework="pt", device="cpu"
        ) as handle:
            for history_chunk in range(chunk):
                source = context_handle if history_chunk == 0 else handle
                history_parts.append(
                    source.get_tensor(f"chunk_{history_chunk:02d}").squeeze(0)
                )
        history_prefeature = torch.cat(history_parts, dim=0).contiguous()

        cross_path = prompt_dir / "cross_attention.safetensors"
        with safe_open(cross_path, framework="pt", device="cpu") as handle:
            cross_k = handle.get_tensor(
                f"block_{self.layer_id:02d}_k"
            ).squeeze(0)
            cross_v = handle.get_tensor(
                f"block_{self.layer_id:02d}_v"
            ).squeeze(0)

        return {
            "prompt_id": prompt_id,
            "chunk": chunk,
            "target_step": target_step,
            "noisy_latent": noisy,
            "timestep": timestep,
            "anchor_timestep": anchor_timestep,
            "anchor_distance": (
                timestep.float() - anchor_timestep.float()
            ).abs().mean(),
            "anchor_hidden": anchor_hidden,
            "previous_hidden": previous_hidden,
            "target_hidden": target_hidden,
            "target_flow": target_flow,
            "history_prefeature": history_prefeature,
            "cross_k": cross_k,
            "cross_v": cross_v,
        }


def collate_lazy_samples(samples: list[dict[str, Any]]) -> dict[str, Any]:
    if not samples:
        raise ValueError("Empty batch")
    chunks = {sample["chunk"] for sample in samples}
    target_steps = {sample["target_step"] for sample in samples}
    if len(chunks) != 1 or len(target_steps) != 1:
        raise ValueError("A batch must share chunk and target_step")
    return {
        "prompt_ids": [sample["prompt_id"] for sample in samples],
        "chunk": samples[0]["chunk"],
        "anchor_step": samples[0]["target_step"] - 1,
        "target_step": samples[0]["target_step"],
        **{
            key: torch.stack([sample[key] for sample in samples])
            for key in (
                "noisy_latent",
                "timestep",
                "anchor_timestep",
                "anchor_distance",
                "anchor_hidden",
                "previous_hidden",
                "target_hidden",
                "target_flow",
                "history_prefeature",
                "cross_k",
                "cross_v",
            )
        },
    }


class ScheduledBatchSampler:
    def __init__(self, batches: list[list[tuple[int, int, int]]]) -> None:
        self.batches = batches

    def __iter__(self):
        return iter(self.batches)

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