# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Small NPZ training adapter using the same PAI contract as the quickstarts."""

import os
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset, DistributedSampler, default_collate

from utils.pai_batch import make_pai_batch


class ProcessedPAI(Dataset):
    def __init__(self, root, frames, tokenizer=False):
        self.paths = sorted(Path(root).rglob("*.npz"))
        if not self.paths:
            raise FileNotFoundError(f"No processed samples under {root}; run prepare_pai_bidir.py first")
        self.frames = frames
        self.tokenizer = tokenizer

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

    def __getitem__(self, index):
        path = self.paths[index]
        with np.load(path, allow_pickle=False) as data:
            sample = {key: data[key] for key in data.files}
        ranges = sample["range_m"]
        if ranges.shape[1:] != (128, 1024) or len(ranges) < self.frames:
            raise ValueError(f"{path}: need {self.frames} real 128x1024 sweeps")
        if (str(sample["dataset"]) != "pandar" or int(sample["sensor_dataset_id"]) != 0
                or float(sample["min_range_m"]) != 2.5 or float(sample["max_range_m"]) != 100
                or float(sample["fps"]) != 10 or bool(sample["has_prepend"])):
            raise ValueError(f"{path}: sample does not match the PAI release contract")
        start = int(torch.randint(len(ranges) - self.frames + 1, ()).item())
        temporal = ("range_m", "valid_mask", "hdmap", "pose_sensor_start_end",
                    "box_valid", "box_object_to_world", "box_lwh", "box_track_id", "box_is_moving")
        for key in temporal:
            if key in sample:
                sample[key] = sample[key][start:start + self.frames]
        if self.tokenizer:
            ranges = sample["range_m"]
            valid = sample["valid_mask"] & np.isfinite(ranges) & (ranges >= 2.5) & (ranges <= 100)
            video = np.where(valid, 2 * (ranges - 2.5) / 97.5 - 1, -1).astype(np.float32)
            # Match tokenizer training's repeated first sweep, not the DiT layout.
            video = torch.from_numpy(np.concatenate((video[:1], video)))[None]
            mask = torch.from_numpy(np.concatenate((valid[:1], valid)))[None].float()
            batch = dict(video=video, recon_target=video.clone(), loss_mask=mask,
                         loss_mask_is_validity=True, sensor_dataset_id=torch.tensor(0),
                         sensor_ray_directions=torch.from_numpy(sample["ray_directions"]).float(),
                         sensor_depth_bounds=torch.tensor([2.5, 100.]))
        else:
            batch = make_pai_batch(sample, self.frames)
            batch = {key: value[0] if torch.is_tensor(value) else value for key, value in batch.items()}
        batch.update(clip_name=path.stem, data_type="lidar", __key__=path.stem)
        return batch


def _collate(samples):
    batch = default_collate(samples)
    for key in ("is_preprocessed", "loss_mask_is_validity"):
        if key in batch:
            batch[key] = True
    return batch


def build_loader(root, frames=21, tokenizer=False, validation=False, num_workers=2, batch_size=1):
    dataset = ProcessedPAI(root, frames, tokenizer)
    sampler = DistributedSampler(dataset, shuffle=not validation)
    return DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers,
                      pin_memory=True, drop_last=False, collate_fn=_collate)


def loader_config(*, tokenizer=False, validation=False):
    from utils.lazy_config import LazyCall as L

    root = os.environ.get("LIDARVERSE_PAI_VAL_DATA" if validation else "LIDARVERSE_PAI_TRAIN_DATA",
                          "data/bidir_quick_start/nvidia_pai")
    return L(build_loader)(root=root, frames=8 if tokenizer else 21, tokenizer=tokenizer,
                           validation=validation, num_workers=2, batch_size=1)