# Copyright 2026 Toyota Research Institute.  All rights reserved.

import argparse
import numpy as np

from anydata.utils.write import write_image
from anydata.dataloaders.instantiate import dataset_from_config
from anydata.utils.render import render_nvidia_world_scenario


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)

    data = dataset[0]
    rendered = render_nvidia_world_scenario(data, has_batch_dim=False)
    for key, val in rendered.items():        
        rgb = (data['rgb'][key].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        write_image(f'rendered/{key}_rgb.png', rgb)
        write_image(f'rendered/{key}_nvidia.png', val)

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