"""ToTensor transformation."""

import numpy as np
import torch

from vis4d.data.const import CommonKeys as K
from vis4d.data.typing import DictData

from .base import Transform


def _replace_arrays(data: DictData) -> None:
    """Replace numpy arrays with tensors."""
    for key in data.keys():
        if key in [K.images, K.original_images]:
            if not data[key].flags.c_contiguous:
                data[key] = np.ascontiguousarray(
                    data[key].transpose(0, 3, 1, 2)
                )
                data[key] = torch.from_numpy(data[key])
            else:
                data[key] = (
                    torch.from_numpy(data[key])
                    .permute(0, 3, 1, 2)
                    .contiguous()
                )
        elif isinstance(data[key], np.ndarray):
            data[key] = torch.from_numpy(data[key])
        elif isinstance(data[key], dict):
            _replace_arrays(data[key])
        elif isinstance(data[key], list):
            for i, entry in enumerate(data[key]):
                if isinstance(entry, np.ndarray):
                    data[key][i] = torch.from_numpy(entry)


@Transform("data", "data")
class ToTensor:
    """Transform all entries in a list of DataDict from numpy to torch.

    Note that we reshape K.images from NHWC to NCHW.
    """

    def __call__(self, batch: list[DictData]) -> list[DictData]:
        """Transform all entries to tensor."""
        for data in batch:
            _replace_arrays(data)
        return batch