# Copyright 2026 Toyota Research Institute.  All rights reserved.

import torch
import torchvision.transforms as transforms
import numpy as np 

from anydata.utils.decorators import iterate1
from anydata.utils.data import make_list


@iterate1
def to_tensor(matrix, tensor_type='torch.FloatTensor'):
    """Casts a matrix to a torch.Tensor"""
    if matrix is None:
        return None
    if isinstance(matrix, dict):
        return {key: to_tensor(val) for key, val in matrix.items()}
    else:
        if isinstance(matrix, np.float32):
            matrix = np.array(matrix)  # Convert to numpy
        if isinstance(matrix, float):
            matrix = np.array([matrix])  # Convert to numpy
        if isinstance(matrix, list) and isinstance(matrix[0], dict):
            return matrix  # NOT HANDLED YET
        if isinstance(matrix, str):
            return matrix  # NOT HANDLED YET
        if matrix.dtype == object:
            return matrix  # NOT HANDLED YET
        return torch.Tensor(matrix).type(tensor_type)


@iterate1
@iterate1
def to_tensor_permute(matrix, tensor_type='torch.FloatTensor'):
    """Casts a matrix to a torch.Tensor"""
    return torch.Tensor(matrix).type(tensor_type).permute(2, 0, 1)


@iterate1
def to_tensor_image(image, tensor_type='torch.FloatTensor'):
    """Casts an image to a torch.Tensor"""
    transform = transforms.ToTensor()
    return transform(image).type(tensor_type)


def to_tensor_sample(sample, tensor_type='torch.FloatTensor', no_to_tensor=False, **kwargs):
    """Casts a sample dictionary to a torch.Tensor"""
    # Do nothing if requested
    if no_to_tensor:
        return sample
    # Convert using torchvision
    keys = ['rgb', 'mask']
    for key_sample, val_sample in sample.items():
        for key in keys:
            if key in key_sample:
                sample[key_sample] = to_tensor_image(val_sample, tensor_type)
    # Convert from numpy
    keys = ['intrinsics', 'extrinsics', 'pose', 'depth', 'semantic', 'timestep', 'action', 'ego']
    for key_sample, val_sample in sample.items():
        for key in keys:
            if key in key_sample:
                sample[key_sample] = to_tensor(val_sample, tensor_type)
    # Convert from numpy while permuting dimension order (012 to 201)
    keys = ['optflow']
    for key_sample, val_sample in sample.items():
        for key in keys:
            if key in key_sample:
                sample[key_sample] = to_tensor_permute(val_sample, tensor_type)
    # Return converted sample
    return sample