import sys

import lovely_tensors
import numpy as np
import torch
from lovely_numpy import lo
from rich import print
from tqdm import tqdm

from custom.dataset.vidar_dataset import VidarDataset

lovely_tensors.monkey_patch()
np.set_printoptions(precision=3, suppress=True)
torch.set_printoptions(precision=3, sci_mode=False, threshold=1000)


# Example: python custom/display/inspect.py ddad

def main_with(argv):

    print(f'====> {argv} <====')

    # dataset_dir = 'custom/display/config/default.yaml'
    # dataset = VidarDataset(dataset_dir=dataset_dir, config_key=sys.argv[1])

    dataset_dir, config_key = argv, 'dataset'
    dataset = VidarDataset(dataset_dir=dataset_dir, config_key=config_key)

    # data = dataset[0]
    data = dataset[10]
    vidar = data['anydata']

    # for i in tqdm(range(len(dataset))):
    #     data = dataset[i]
    #     vidar = data['anydata']
    #     # print(vidar['filename'][(0,0)])
    #     # if len(vidar['rgb']) != 41 * 4:
    #     #     print(i, vidar['filename'], len(vidar['rgb']))

    filename = vidar['filename']
    if isinstance(filename, dict):
        filename = filename[(0,0)]
    tag = vidar['tag']

    labels = [key for key in vidar.keys()]
    labels = str(labels).replace("'", "").replace(' ', '')
    time = list(set([key[0] for key in vidar['time_cam'].keys()]))
    cam  = list(set([key[1] for key in vidar['time_cam'].keys()]))

    has_ego = any([key[1] == 'ego' for key in vidar['pose'].keys()])

    if 'action' in vidar.keys():
        action = vidar['action'][(0,0)]
        if isinstance(action, torch.Tensor):
            action = action.numpy().shape
        elif isinstance(action, dict):
            action = action['action']
            if isinstance(action, dict):
                action = action.keys()
            else:
                action = action.shape
    else:
        action = None

    print(f'######################################### VIDAR')
    print(f'### Filename: {filename}')
    print(f'### Tag: {tag}')
    print(f'### Labels: {labels}')
    print(f'### Timesteps: {len(time)}')
    print(f'### Cameras ({len(cam)}): {cam}')
    print(f'### Has Ego: {has_ego}')
    if action is not None:
        print(f'### Action: {action}')
    print(f'######################################### COSMOS')
    for key in ['dset_name', 'prompt', 'fps', 'scale_factor', 'image_size']:
        print(f'### {key}: {data[key]}')
    print(f'#########################################')
    print(f'### Length: {len(dataset)}')
    print(f'#########################################')


if __name__ == "__main__":

    for argv in sys.argv[1:]:
        main_with(argv)