# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import yaml
import importlib

from copy import deepcopy
from functools import partial 
from anydata.utils.read import read_config
from anydata.augmentations.crop import crop_sample
from anydata.augmentations.resize import resize_sample
from anydata.augmentations.tensor import to_tensor_sample
from anydata.augmentations.misc import depth_filter


def dataset_from_config(config, name='dataset', added=None, subfolder=None, only_camera=None):
    # Read configuration and dataset name
    if isinstance(config, str):
        config: dict = read_config(config)[name]
    name = config['name']

    # Look into a subfolder of the dataset insterad (e.g., LBM task)
    if subfolder is not None:
        for i in range(len(config['path'])):
            split = config['path'][i].split('/')
            config['path'][i] = "/" + os.path.join(*split[:-1], subfolder, *split[-1:])
    
    # Extract custom metadata (stored in dataset, added to batch metadata)
    custom = config.pop('custom', None) or config.pop('added', None)  # 'added' is legacy name
    config['custom'] = custom

    # Remap legacy config key names to match renamed constructor params
    # NOTE(bvh): webdataset configs still contain cameras (= the only key,
    # since other keys are not relevant for them), so don't remove this code
    _KEY_REMAP = {
        'cameras': 'cameras_core',
        'cameras_sample': 'cameras_core_sample',
        'cameras_context': 'cameras_aux',
        'cameras_context_sample': 'cameras_aux_sample',
    }
    for old_key, new_key in _KEY_REMAP.items():
        if old_key in config and new_key not in config:
            config[new_key] = config.pop(old_key)

    # Include provided additional keys
    if added is not None:
        for key, val in added.items():
            config[key] = val

    # Prepare augmentation
    config_augmentation = config.get('augmentation', None)
    if config_augmentation is None:
        config['data_transform'] = to_tensor_sample
    else:
        config['data_transform'] = partial(augmentations, cfg=config['augmentation'])

    if only_camera is not None:
        config['cameras'] = [only_camera]
        config['filter_camera'] = only_camera

    # Instantiate and return dataset
    module = importlib.import_module(f'anydata.dataloaders.{name}')
    method = getattr(module, name + 'Dataset', None)
    return method(**config), config


def augmentations(sample, cfg, no_to_tensor=False, webbed=False):
    """Data augmentations"""
    # Loop over augmentation keys
    for key in cfg.keys():

        ### RESIZE
        if key.startswith('resize'):
            shape = cfg[key]
            
            if 'resolution_original' not in sample['metadata']:
                resolution = sample['metadata']['resolution']
                sample['metadata']['resolution_original'] = {}
                
                for cam_idx, cam_name in enumerate(sample['metadata']['cameras']):
                    if isinstance(resolution, dict):
                        # Resolution may be keyed by camera name (str) or camera index (int)
                        sample['metadata']['resolution_original'][cam_idx] = \
                            resolution.get(cam_name, resolution.get(cam_idx, next(iter(resolution.values()))))
                    else:
                        sample['metadata']['resolution_original'][cam_idx] = resolution
            
            # preserve depth if requested or if depth map is sparse
            preserve_depth = cfg.get('preserve_depth', 'depth' in sample['metadata'].get('conventions', {}) and sample['metadata']['conventions']['depth']['sparse'])
            sample, shape = resize_sample(sample, shape=shape, preserve_depth=preserve_depth)
            shape = {cam_idx: val for (t_idx, cam_idx), val in shape.items() if t_idx == 0}  # per-camera shapes only
            
            if 'augmentations' not in sample['metadata']:
                sample['metadata']['augmentations'] = {}
            
            for shape_key, shape_val in shape.items():
                sample['metadata']['augmentations'].setdefault(shape_key, []).append([key, shape_val])
            
            if webbed:  # Webbed data does not take camera names
                sample['metadata']['resolution'] = {
                    cam_idx: val for cam_idx, val in shape.items()}
            else:  # Non-webbed data takes camera names (strings) back
                sample['metadata']['resolution'] = {
                    sample['metadata']['cameras'][cam_idx]: val for cam_idx, val in shape.items()}

        ### CROP
        if key.startswith('crop'):
            shape = cfg[key]

            if 'resolution_original' not in sample['metadata']:
                resolution = sample['metadata']['resolution']
                sample['metadata']['resolution_original'] = {}

                for cam_idx, cam_name in enumerate(sample['metadata']['cameras']):
                    if isinstance(resolution, dict):
                        # Resolution may be keyed by camera name (str) or camera index (int)
                        sample['metadata']['resolution_original'][cam_idx] = \
                            resolution.get(cam_name, resolution.get(cam_idx, next(iter(resolution.values()))))
                    else:
                        sample['metadata']['resolution_original'][cam_idx] = resolution

            # preserve depth if requested or if depth map is sparse
            sample, shape = crop_sample(sample, shape=shape, mode=key)
            shape = {cam_idx: val for (t_idx, cam_idx), val in shape.items() if t_idx == 0}  # per-camera shapes only

            if 'augmentations' not in sample['metadata']:
                sample['metadata']['augmentations'] = {}

            for shape_key, shape_val in shape.items():
                sample['metadata']['augmentations'].setdefault(shape_key, []).append([key, shape_val])

            if webbed:  # Webbed data does not take camera names
                sample['metadata']['resolution'] = {
                    cam_idx: val for cam_idx, val in shape.items()}
            else:  # Non-webbed data takes camera names (strings) back
                sample['metadata']['resolution'] = {
                    sample['metadata']['cameras'][cam_idx]: val for cam_idx, val in shape.items()}

        ### DEPTH AUGMENTATIONS
        if 'depth' in sample:
            if key.startswith('depth_filter'):
                sample['depth'] = depth_filter(
                    sample['depth'], sample['intrinsics'], cfg[key])

    ### TO TENSOR
    if not no_to_tensor:
        sample = to_tensor_sample(sample)
    
    # Return augmented sample
    return sample