"""Test: iterate the full train dataloader with workers, report None-label batches."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def main():
    from src import training
    from src.common import load_jsonl
    rows = load_jsonl(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "manifests", "split_train.jsonl"))
    for seed in (0, 1, 2):
        ds, dl = training.make_dataloader(rows, (96, 96, 96), 3, True, num_workers=4, seed=seed)
        n_none, n = 0, 0
        for i, (img, lab, w) in enumerate(dl):
            if lab is None:
                n_none += 1
                print("  seed", seed, "batch", i, "None label", flush=True)
            n += 1
        print("seed", seed, "batches", n, "none-label", n_none, flush=True)
    print("TEST DONE")


if __name__ == "__main__":
    main()