import dataclasses
import os
import time

from openpi.training import config as c
from openpi.training import data_loader as dl


def main() -> None:
    workers = int(os.environ.get("KINOVA_BENCH_WORKERS", "12"))
    batch_size = int(os.environ.get("KINOVA_BENCH_BATCH_SIZE", "24"))
    cfg = dataclasses.replace(
        c.get_config("pi05_kinova_teddybear"),
        exp_name=f"bench_first_b{batch_size}_w{workers}",
        batch_size=batch_size,
        num_workers=workers,
    )
    print({"workers": workers, "batch_size": batch_size}, flush=True)
    t0 = time.time()
    loader = dl.create_data_loader(cfg, shuffle=False, num_batches=1)
    print({"loader_created_s": round(time.time() - t0, 3)}, flush=True)
    t1 = time.time()
    _ = next(iter(loader))
    print({"first_batch_s": round(time.time() - t1, 3)}, flush=True)


if __name__ == "__main__":
    main()