import dataclasses
import time

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


def main() -> None:
    for workers in [16, 32, 48]:
        cfg = dataclasses.replace(
            c.get_config("pi05_kinova_teddybear"),
            exp_name=f"bench_{workers}",
            num_workers=workers,
        )
        t0 = time.time()
        loader = dl.create_data_loader(cfg, shuffle=True, num_batches=6, skip_norm_stats=True)
        loader_s = time.time() - t0

        it = iter(loader)
        t1 = time.time()
        first_batch_s = None
        for i in range(6):
            _ = next(it)
            if i == 0:
                first_batch_s = time.time() - t1
        total_s = time.time() - t1
        print(
            {
                "workers": workers,
                "loader_s": round(loader_s, 3),
                "first_batch_s": round(first_batch_s or 0.0, 3),
                "avg_batch_s": round(total_s / 6, 3),
            },
            flush=True,
        )


if __name__ == "__main__":
    main()