# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import sys
import torch
import argparse
import importlib
import numpy as np

from PIL import Image

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)
    args = parser.parse_args()

    return args


def main_with(config):

    # 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)

    # Get first sample
    data = dataset[0]

    print('##########################')
    print('##### BATCH')
    print('##########################')
    dict_keys, dict_vals = dict(), dict()
    for key1, val1 in data.items():
        if key1 in ['metadata']:
            continue
        if isinstance(val1, dict):

            fkey2 = list(val1.keys())[0]
            dict_keys[key1] = list(val1.keys())
            dict_vals[key1] = list(val1.values())

            if isinstance(val1[fkey2], torch.Tensor):

                num_cameras = max([time_cam[1] for time_cam in dict_keys[key1]]) + 1
                num_frames = {cam: len([time_cam[0] for time_cam in dict_keys[key1] if time_cam[1] == cam]) for cam in range(num_cameras)}
                resolution = dict_vals[key1][0].shape

                same_frames = all([val == num_frames[0] for val in num_frames.values()])
                if same_frames: num_frames = list([v for v in num_frames.values()])[0]
                same_resolution = all([val.shape == dict_vals[key1][0].shape for val in dict_vals[key1]])

                passed = same_resolution & same_frames
                passed = 'SUCCESS' if passed else 'ERROR'
                print(f'### {key1}:', f'{num_frames}/{num_cameras} - {resolution} - {passed}')

            elif isinstance(val1[fkey2], dict):
                if len(val1[fkey2]) == 0:
                    continue
                fkey3 = list(val1[fkey2].keys())[0]
                print(f'### {key1}:', f'{len(dict_keys[key1])}')
                for key3 in val1[fkey2].keys():

                    if isinstance(val1[fkey2][key3], torch.Tensor) or isinstance(val1[fkey2][key3], np.ndarray):

                        num_frames = len([val[key3] for val in val1.values()])
                        resolution = val1[fkey2][fkey3].shape

                        same_resolution = all([val[key3].shape == val1[fkey2][key3].shape for val in val1.values()])
                        
                        passed = same_resolution
                        passed = 'SUCCESS' if passed else 'ERROR'
                        print(f'###### {key3}:', f'{num_frames} - {resolution} - {passed}')

                    elif isinstance(val1[fkey2][fkey3], str):

                        num_frames = len([val[key3] for val in val1.values()])
                        value = val1[fkey2][key3]

                        print(f'###### {key3}:', f'{num_frames} - {value}')

            else:
                print(f'### {key1}:', {key2: val2 for key2, val2 in val1.items()})
        else:
            print(f'### {key1}: {val1}')
    print('##########################')
    print('##### METADATA')
    print('##########################')
    for key1, val1 in data['metadata'].items():
        if key1 == 'specific':
            continue
        if isinstance(val1, dict):
            print('###', key1)
            for key2, val2 in val1.items():
                if isinstance(val2, dict):
                    print(f'###### {key2}:')
                    for key3, val3 in val2.items():
                        print(f'######### {key3}: {val3}')
                else:
                    print(f'###### {key2}: {val2}')
        else:
            print(f'### {key1}: {val1}')
    print('##########################')
    print('##### CHECKS')
    print('##########################')

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