# Modified in the ReTime research artifact; see NOTICE.
import json
import os
from accelerate.logging import get_logger
import numpy as np
from torch.utils.data import DataLoader
import torch.distributed as dist
from pathlib import Path
from starVLA.dataloader.infoflow_sampler import StatelessDrawSampler

logger = get_logger(__name__)

def save_dataset_statistics(dataset_statistics, run_dir):
    """Saves a `dataset_statistics.json` file."""
    out_path = run_dir / "dataset_statistics.json"
    with open(out_path, "w") as f_json:
        for _, stats in dataset_statistics.items():
            for k in stats["action"].keys():
                if isinstance(stats["action"][k], np.ndarray):
                    stats["action"][k] = stats["action"][k].tolist()
            if "proprio" in stats:
                for k in stats["proprio"].keys():
                    if isinstance(stats["proprio"][k], np.ndarray):
                        stats["proprio"][k] = stats["proprio"][k].tolist()
            if "num_trajectories" in stats:
                if isinstance(stats["num_trajectories"], np.ndarray):
                    stats["num_trajectories"] = stats["num_trajectories"].item()
            if "num_transitions" in stats:
                if isinstance(stats["num_transitions"], np.ndarray):
                    stats["num_transitions"] = stats["num_transitions"].item()
        json.dump(dataset_statistics, f_json, indent=2)
    logger.info(f"Saved dataset statistics file at path {out_path}")



def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"):

    if dataset_py == "lerobot_datasets":
        from starVLA.dataloader.lerobot_datasets import get_vla_dataset, collate_fn
        vla_dataset_cfg = cfg.datasets.vla_data

        vla_dataset = get_vla_dataset(
            data_cfg=vla_dataset_cfg,
            balance_dataset_weights=vla_dataset_cfg.get("balance_dataset_weights", False),
            balance_trajectory_weights=vla_dataset_cfg.get("balance_trajectory_weights", False),
        )

        traceable_sampling = bool(vla_dataset_cfg.get("traceable_sampling", False))
        per_device_batch_size = int(cfg.datasets.vla_data.per_device_batch_size)
        world_size = (
            dist.get_world_size()
            if dist.is_available() and dist.is_initialized()
            else int(vla_dataset_cfg.get("expected_world_size", 1))
        )
        sampler = None
        if traceable_sampling:
            expected_world_size = int(vla_dataset_cfg.expected_world_size)
            if world_size != expected_world_size:
                raise RuntimeError(
                    f"Traceable sampling requires world size {expected_world_size}, "
                    f"but distributed runtime reports {world_size}"
                )
            gradient_accumulation_steps = int(
                cfg.trainer.gradient_accumulation_steps
            )
            global_update_batch = (
                per_device_batch_size
                * world_size
                * gradient_accumulation_steps
            )
            usable_draws = len(vla_dataset) - (
                len(vla_dataset) % global_update_batch
            )
            if usable_draws <= 0:
                raise ValueError(
                    f"Dataset length {len(vla_dataset)} is smaller than one global "
                    f"optimizer update ({global_update_batch} draws)"
                )
            sampler = StatelessDrawSampler(usable_draws)
        num_workers = int(vla_dataset_cfg.get("num_workers", 16))
        dataloader_kwargs = dict(
            dataset=vla_dataset,
            batch_size=per_device_batch_size,
            collate_fn=collate_fn,
            sampler=sampler,
            num_workers=num_workers,
            pin_memory=True,
            persistent_workers=num_workers > 0,
            drop_last=traceable_sampling,
        )
        if num_workers > 0:
            dataloader_kwargs["prefetch_factor"] = int(vla_dataset_cfg.get("prefetch_factor", 4))
        vla_train_dataloader = DataLoader(**dataloader_kwargs)

        rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
        if rank == 0:
            output_dir = Path(cfg.output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)
            vla_dataset.save_dataset_statistics(output_dir / "dataset_statistics.json")
            if vla_dataset.valid_anchor_manifest is not None:
                manifest_path = output_dir / "valid_anchor_manifest.json"
                vla_dataset.valid_anchor_manifest.save(manifest_path)
        return vla_train_dataloader
    raise ValueError(f"Unsupported dataset loader: {dataset_py}")