# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import torch
import numpy as np
from typing import Any, Dict, List, Optional, Union

from anydata.utils.types import is_seq, is_tensor, is_dict


def get_local_root():
    """Return local data root from $ANYDATA_LOCAL_ROOT env var (default: /data)."""
    return os.environ.get('ANYDATA_LOCAL_ROOT', '/data')


def frame_name(i, add=0):
    """Name for high-dimensional label"""
    if isinstance(i, str):
        i = os.path.basename(i).split('.')[0]
        if '_' in i: i = i.split('_')[-1]
        i = i.replace('DSCF', '').replace('DSC', '').replace('_', '')
        i = int(i)        
    return '%010d' % (i + add)


def update_dict(data, key1, key2, val):
    """Update dict with a new key"""
    # Return None if it is None already
    if val is None:
        return
    # Create first key if it does not exist
    if key1 not in data:
        data[key1] = {}
    # Create second key if it does not exist
    if key2 not in data[key1]:
        data[key1][key2] = {}
    # Assign second key
    data[key1][key2] = val


def update_dict_nested(data, key1, key2, key3, val):
    """Update dict of dicts with a new key"""
    # Return None if it is None already
    if val is None:
        return
    # Create first key if it does not exist
    if key1 not in data:
        data[key1] = {}
    # Create second key if it does not exist
    if key2 not in data[key1]:
        data[key1][key2] = {}
    # Create third key if it does not exist
    if key3 not in data[key1][key2]:
        data[key1][key2][key3] = {}
    # Assign third key
    data[key1][key2][key3] = val


def stack_sample(sample: List[Dict[str, Any]], lidar_sample=None, radar_sample=None):
    """Stack sample key by key, along the batch dimension"""
    # If there are no tensors, return empty list
    if len(sample) == 0:
        return None
    # If there is only one sensor don't do anything
    if len(sample) == 1:
        sample = sample[0]
        return sample
    # Otherwise, stack sample
    first_sample = sample[0]
    stacked_sample: Dict[str, list | torch.Tensor | np.ndarray | dict | None] = {}
    for key, val in first_sample.items():
        # Global keys (do not stack)
        if key in ['idx', 'dataset_idx']:
            stacked_sample[key] = val
        # Meta keys
        elif key in ['meta']:
            stacked_sample[key] = {}
            for key2 in val.keys():
                stacked_sample[key][key2] = {}
                for key3 in val[key2].keys():
                    stacked_sample[key][key2][key3] = torch.stack(
                        [torch.tensor(s[key][key2][key3]) for s in sample], 0)
        # Stack numpy arrays
        elif isinstance(val, np.ndarray):
            array_list: List[np.ndarray] = [s[key] for s in sample if key in s]
            homogeneous = all(arr.shape == val.shape for arr in array_list)
            if homogeneous:
                stacked_sample[key] = np.stack(array_list, 0)
            else:  # Others can be stacked
                stacked_sample[key] = np.array(array_list, dtype=np.object_)
        # Stack tensors
        elif is_tensor(val):
            stacked_sample[key] = torch.stack([s[key] for s in sample], 0)
        # Stack list
        elif is_seq(val):
            stacked_sample[key] = []
            # Stack list of torch tensors
            if len(val) > 0 and is_tensor(val[0]):
                for i in range(len(val)):
                    stacked_sample[key].append(torch.stack([s[key][i] for s in sample], 0))
            else:
                stacked_sample[key] = [s[key] for s in sample]
        # Repeat for dictionaries
        elif is_dict(val):
            stacked_sample[key] = stack_sample([s[key] for s in sample])
        elif isinstance(val, str) or val is None:
            stacked_sample[key] = np.stack([s[key] for s in sample])
        # Append lists
        else:
            stacked_sample[key] = [s[key] for s in sample]

    # Return stacked sample
    return stacked_sample


def join_dict(dct_or_item: Union[dict, Any], prefix: Optional[str] = None, delimiter: str = "/"):
    """
    Collapses a nested dict into a single level dict where nested dicts are combined with the
    parent key by appending the nested dict key separated by the given delimiter.
    """
    if isinstance(dct_or_item, dict):
        new_dict = {}
        for key, val in dct_or_item.items():
            sub_prefix = f"{prefix}{delimiter}{key}" if prefix is not None else key
            val = join_dict(val, prefix=sub_prefix, delimiter=delimiter)
            new_dict.update(val)
        return new_dict
    return {prefix: dct_or_item}


def unjoin_dict(dct: Dict[str, Any], delimiter: str = "/"):
    """
    Undoes the effect of the join_dict function by splitting a single key into dicts for each
    delimiter in the key
    """
    unjoined_dict = {}
    for k, v in dct.items():
        keylist = k.split(delimiter)
        subdict = unjoined_dict
        for subkey in keylist[:-1]:
            if subkey not in subdict:
                subdict[subkey] = {}
            subdict = subdict[subkey]
        subdict[keylist[-1]] = v
    return unjoined_dict