# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import argparse

from anydata.dataloaders.instantiate import dataset_from_config


def parse_args(num_procs=24):
    parser = argparse.ArgumentParser()
    parser.add_argument('config', type=str)
    parser.add_argument('--subset', type=str, default=None)
    parser.add_argument(
        '--visualizer', type=str, default='camviz',
        choices=['camviz', 'rerun'],
        help='Visualizer backend: "camviz" (default) or "rerun"',
    )
    parser.add_argument(
        '--stride', type=int, default=1,
        help='Step between visualized samples (rerun only; default: 1)',
    )
    args = parser.parse_args()

    return args


def main_with(args):

    # If a subdataset is provided, parse into config and name
    if '/' not in args.config:
        config, name = 'anydata/inspect/unified.yaml', args.config
    elif not args.config.endswith('yaml'):
        config, name = os.path.dirname(args.config), os.path.basename(args.config)
    # Not provide, provide only config
    else: config, name = args.config, None

    added = {}
    if args.subset is not None:
        added['subset'] = args.subset

    # Instantiate dataset
    dataset, _ = dataset_from_config(config, name, added)

    # Display dataset with the chosen visualizer
    visualizer = getattr(args, 'visualizer', 'camviz')
    if visualizer == 'rerun':
        from anydata.inspect.rerun_display import RerunDisplay
        display = RerunDisplay(dataset)
        loop_kwargs = {'stride': args.stride}
    else:
        from anydata.inspect.camviz_display import CamvizDisplay
        display = CamvizDisplay(dataset)
        loop_kwargs = {}
    display.loop(**loop_kwargs)


if __name__ == "__main__":
    args = parse_args()
    main_with(args)