VideoTensor

@classmethod
def __docarray_validate_getitem__(cls, item: Any) -> Tuple[int]:
    """This method validates the input to `AbstractTensor.__class_getitem__`.

    It is called at "class creation time",
    i.e. when a class is created with syntax of the form AnyTensor[shape].

    The default implementation tries to cast any `item` to a tuple of ints.
    A subclass can override this method to implement custom validation logic.

    The output of this is eventually passed to
    [`AbstractTensor.__docarray_validate_shape__`]
    [docarray.typing.tensor.abstract_tensor.AbstractTensor.__docarray_validate_shape__]
    as its `shape` argument.

    Raises `ValueError` if the input `item` does not pass validation.

    :param item: The item to validate, passed to `__class_getitem__` (`Tensor[item]`).
    :return: The validated item == the target shape of this tensor.
    """
    if isinstance(item, int):
        item = (item,)
    try:
        item = tuple(item)
    except TypeError:
        raise TypeError(f'{item} is not a valid tensor shape.')
    return item

@classmethod
def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str], ...]) -> T:
    """Every tensor has to implement this method in order to
    enable syntax of the form AnyTensor[shape].
    It is called when a tensor is assigned to a field of this type.
    i.e. when a tensor is passed to a Document field of type AnyTensor[shape].

    The intended behaviour is as follows:

    - If the shape of `t` is equal to `shape`, return `t`.
    - If the shape of `t` is not equal to `shape`,
        but can be reshaped to `shape`, return `t` reshaped to `shape`.
    - If the shape of `t` is not equal to `shape`
        and cannot be reshaped to `shape`, raise a ValueError.

    :param t: The tensor to validate.
    :param shape: The shape to validate against.
    :return: The validated tensor.
    """
    comp_be = t.get_comp_backend()
    tshape = comp_be.shape(t)
    if tshape == shape:
        return t
    elif any(isinstance(dim, str) or dim == Ellipsis for dim in shape):
        ellipsis_occurrences = [
            pos for pos, dim in enumerate(shape) if dim == Ellipsis
        ]
        if ellipsis_occurrences:
            if len(ellipsis_occurrences) > 1:
                raise ValueError(
                    f'Cannot use Ellipsis (...) more than once for the shape {shape}'
                )
            ellipsis_pos = ellipsis_occurrences[0]
            # Calculate how many dimensions to add. Should be at least 1.
            dimensions_needed = max(len(tshape) - len(shape) + 1, 1)
            shape = (
                shape[:ellipsis_pos]
                + tuple(
                    f'__dim_var_{index}__' for index in range(dimensions_needed)
                )
                + shape[ellipsis_pos + 1 :]
            )

        if len(tshape) != len(shape):
            raise ValueError(
                f'Tensor shape mismatch. Expected {shape}, got {tshape}'
            )
        known_dims: Dict[str, int] = {}
        for tdim, dim in zip(tshape, shape):
            if isinstance(dim, int) and tdim != dim:
                raise ValueError(
                    f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                )
            elif isinstance(dim, str):
                if dim in known_dims and known_dims[dim] != tdim:
                    raise ValueError(
                        f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                    )
                else:
                    known_dims[dim] = tdim
        else:
            return t
    else:
        shape = cast(Tuple[int], shape)
        warnings.warn(
            f'Tensor shape mismatch. Reshaping tensor '
            f'of shape {tshape} to shape {shape}'
        )
        try:
            value = cls._docarray_from_native(comp_be.reshape(t, shape))
            return cast(T, value)
        except RuntimeError:
            raise ValueError(
                f'Cannot reshape tensor of shape {tshape} to shape {shape}'
            )

@abc.abstractmethod
def __getitem__(self: T, item) -> T:
    """Get a slice of this tensor."""
    ...

@abc.abstractmethod
def __iter__(self):
    """Iterate over the elements of this tensor."""
    ...

@abc.abstractmethod
def __setitem__(self, index, value):
    """Set a slice of this tensor."""
    ...

def display(self, audio: Optional[AudioTensor] = None) -> None:
    """
    Display video data from tensor in notebook.

    :param audio: sound to play with video tensor
    """
    if is_notebook():
        from IPython.display import Video, display

        b = self.to_bytes(audio_tensor=audio)
        display(Video(data=b, embed=True, mimetype='video/mp4'))
    else:
        warnings.warn('Display of video is only possible in a notebook.')

@classmethod
def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T':
    """
    Read ndarray from a proto msg
    :param pb_msg:
    :return: a numpy array
    """
    source = pb_msg.dense
    if source.buffer:
        x = np.frombuffer(bytearray(source.buffer), dtype=source.dtype)
        return cls._docarray_from_native(x.reshape(source.shape))
    elif len(source.shape) > 0:
        return cls._docarray_from_native(np.zeros(source.shape))
    else:
        raise ValueError(f'proto message {pb_msg} cannot be cast to a NdArray')

@staticmethod
def get_comp_backend() -> 'NumpyCompBackend':
    """Return the computational backend of the tensor"""
    from docarray.computation.numpy_backend import NumpyCompBackend

    return NumpyCompBackend()

def save(
    self: 'T',
    file_path: Union[str, BytesIO],
    audio_tensor: Optional[AudioTensor] = None,
    video_frame_rate: int = 24,
    video_codec: str = 'h264',
    audio_frame_rate: int = 48000,
    audio_codec: str = 'aac',
    audio_format: str = 'fltp',
) -> None:
    """
    Save video tensor to a .mp4 file.

    ---

    ```python
    import numpy as np

    from docarray import BaseDoc
    from docarray.typing.tensor.audio.audio_tensor import AudioTensor
    from docarray.typing.tensor.video.video_tensor import VideoTensor


    class MyDoc(BaseDoc):
        video_tensor: VideoTensor
        audio_tensor: AudioTensor


    doc = MyDoc(
        video_tensor=np.random.randint(low=0, high=256, size=(10, 200, 300, 3)),
        audio_tensor=np.random.randn(100, 1, 1024).astype("float32"),
    )

    doc.video_tensor.save(
        file_path="/tmp/mp_.mp4",
        audio_tensor=doc.audio_tensor,
        audio_format="flt",
    )
    ```

    ---
    :param file_path: path to a .mp4 file. If file is a string, open the file by
        that name, otherwise treat it as a file-like object.
    :param audio_tensor: AudioTensor containing the video's soundtrack.
    :param video_frame_rate: video frames per second.
    :param video_codec: the name of a video decoder/encoder.
    :param audio_frame_rate: audio frames per second.
    :param audio_codec: the name of an audio decoder/encoder.
    :param audio_format: the name of one of the audio formats supported by PyAV,
        such as 'flt', 'fltp', 's16' or 's16p'.
    """
    if TYPE_CHECKING:
        import av
    else:
        av = import_library('av', raise_error=True)

    np_tensor = self.get_comp_backend().to_numpy(array=self)
    video_tensor = np_tensor.astype('uint8')

    if isinstance(file_path, str):
        format = file_path.split('.')[-1]
    else:
        format = 'mp4'

    with av.open(file_path, mode='w', format=format) as container:
        if video_tensor.ndim == 3:
            video_tensor = np.expand_dims(video_tensor, axis=0)

        stream_video = container.add_stream(video_codec, rate=video_frame_rate)
        stream_video.height = video_tensor.shape[-3]
        stream_video.width = video_tensor.shape[-2]

        if audio_tensor is not None:
            stream_audio = container.add_stream(audio_codec)
            audio_np = audio_tensor.get_comp_backend().to_numpy(array=audio_tensor)
            audio_layout = 'stereo' if audio_np.shape[-2] == 2 else 'mono'

            for i, audio in enumerate(audio_np):
                frame = av.AudioFrame.from_ndarray(
                    array=audio, format=audio_format, layout=audio_layout
                )
                frame.rate = audio_frame_rate
                frame.pts = audio.shape[-1] * i
                for packet in stream_audio.encode(frame):
                    container.mux(packet)

            for packet in stream_audio.encode(None):
                container.mux(packet)

        for vid in video_tensor:
            frame = av.VideoFrame.from_ndarray(vid, format='rgb24')
            for packet in stream_video.encode(frame):
                container.mux(packet)

        for packet in stream_video.encode(None):
            container.mux(packet)

def to_bytes(
    self: 'T',
    audio_tensor: Optional[AudioTensor] = None,
    video_frame_rate: int = 24,
    video_codec: str = 'h264',
    audio_frame_rate: int = 48000,
    audio_codec: str = 'aac',
    audio_format: str = 'fltp',
) -> 'VideoBytes':
    """
    Convert video tensor to [`VideoBytes`][docarray.typing.VideoBytes].

    :param audio_tensor: AudioTensor containing the video's soundtrack.
    :param video_frame_rate: video frames per second.
    :param video_codec: the name of a video decoder/encoder.
    :param audio_frame_rate: audio frames per second.
    :param audio_codec: the name of an audio decoder/encoder.
    :param audio_format: the name of one of the audio formats supported by PyAV,
        such as 'flt', 'fltp', 's16' or 's16p'.

    :return: a VideoBytes object
    """
    from docarray.typing.bytes.video_bytes import VideoBytes

    bytes = BytesIO()
    self.save(
        file_path=bytes,
        audio_tensor=audio_tensor,
        video_frame_rate=video_frame_rate,
        video_codec=video_codec,
        audio_frame_rate=audio_frame_rate,
        audio_codec=audio_codec,
        audio_format=audio_format,
    )
    return VideoBytes(bytes.getvalue())

def to_protobuf(self) -> 'NdArrayProto':
    """
    Transform self into a NdArrayProto protobuf message
    """
    from docarray.proto import NdArrayProto

    nd_proto = NdArrayProto()

    nd_proto.dense.buffer = self.tobytes()
    nd_proto.dense.ClearField('shape')
    nd_proto.dense.shape.extend(list(self.shape))
    nd_proto.dense.dtype = self.dtype.str

    return nd_proto

def unwrap(self) -> np.ndarray:
    """
    Return the original ndarray without any memory copy.

    The original view rest intact and is still a Document `NdArray`
    but the return object is a pure `np.ndarray` but both object share
    the same memory layout.

    ---

    ```python
    from docarray.typing import NdArray
    import numpy as np
    from pydantic import parse_obj_as

    t1 = parse_obj_as(NdArray, np.zeros((3, 224, 224)))
    t2 = t1.unwrap()
    # here t2 is a pure np.ndarray but t1 is still a Docarray NdArray
    # But both share the same underlying memory
    ```

    ---

    :return: a `numpy.ndarray`
    """
    return self.view(np.ndarray)

@classmethod
def __docarray_validate_getitem__(cls, item: Any) -> Tuple[int]:
    """This method validates the input to `AbstractTensor.__class_getitem__`.

    It is called at "class creation time",
    i.e. when a class is created with syntax of the form AnyTensor[shape].

    The default implementation tries to cast any `item` to a tuple of ints.
    A subclass can override this method to implement custom validation logic.

    The output of this is eventually passed to
    [`AbstractTensor.__docarray_validate_shape__`]
    [docarray.typing.tensor.abstract_tensor.AbstractTensor.__docarray_validate_shape__]
    as its `shape` argument.

    Raises `ValueError` if the input `item` does not pass validation.

    :param item: The item to validate, passed to `__class_getitem__` (`Tensor[item]`).
    :return: The validated item == the target shape of this tensor.
    """
    if isinstance(item, int):
        item = (item,)
    try:
        item = tuple(item)
    except TypeError:
        raise TypeError(f'{item} is not a valid tensor shape.')
    return item

@classmethod
def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str], ...]) -> T:
    """Every tensor has to implement this method in order to
    enable syntax of the form AnyTensor[shape].
    It is called when a tensor is assigned to a field of this type.
    i.e. when a tensor is passed to a Document field of type AnyTensor[shape].

    The intended behaviour is as follows:

    - If the shape of `t` is equal to `shape`, return `t`.
    - If the shape of `t` is not equal to `shape`,
        but can be reshaped to `shape`, return `t` reshaped to `shape`.
    - If the shape of `t` is not equal to `shape`
        and cannot be reshaped to `shape`, raise a ValueError.

    :param t: The tensor to validate.
    :param shape: The shape to validate against.
    :return: The validated tensor.
    """
    comp_be = t.get_comp_backend()
    tshape = comp_be.shape(t)
    if tshape == shape:
        return t
    elif any(isinstance(dim, str) or dim == Ellipsis for dim in shape):
        ellipsis_occurrences = [
            pos for pos, dim in enumerate(shape) if dim == Ellipsis
        ]
        if ellipsis_occurrences:
            if len(ellipsis_occurrences) > 1:
                raise ValueError(
                    f'Cannot use Ellipsis (...) more than once for the shape {shape}'
                )
            ellipsis_pos = ellipsis_occurrences[0]
            # Calculate how many dimensions to add. Should be at least 1.
            dimensions_needed = max(len(tshape) - len(shape) + 1, 1)
            shape = (
                shape[:ellipsis_pos]
                + tuple(
                    f'__dim_var_{index}__' for index in range(dimensions_needed)
                )
                + shape[ellipsis_pos + 1 :]
            )

        if len(tshape) != len(shape):
            raise ValueError(
                f'Tensor shape mismatch. Expected {shape}, got {tshape}'
            )
        known_dims: Dict[str, int] = {}
        for tdim, dim in zip(tshape, shape):
            if isinstance(dim, int) and tdim != dim:
                raise ValueError(
                    f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                )
            elif isinstance(dim, str):
                if dim in known_dims and known_dims[dim] != tdim:
                    raise ValueError(
                        f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                    )
                else:
                    known_dims[dim] = tdim
        else:
            return t
    else:
        shape = cast(Tuple[int], shape)
        warnings.warn(
            f'Tensor shape mismatch. Reshaping tensor '
            f'of shape {tshape} to shape {shape}'
        )
        try:
            value = cls._docarray_from_native(comp_be.reshape(t, shape))
            return cast(T, value)
        except RuntimeError:
            raise ValueError(
                f'Cannot reshape tensor of shape {tshape} to shape {shape}'
            )

@abc.abstractmethod
def __getitem__(self: T, item) -> T:
    """Get a slice of this tensor."""
    ...

@abc.abstractmethod
def __iter__(self):
    """Iterate over the elements of this tensor."""
    ...

@abc.abstractmethod
def __setitem__(self, index, value):
    """Set a slice of this tensor."""
    ...

def display(self, audio: Optional[AudioTensor] = None) -> None:
    """
    Display video data from tensor in notebook.

    :param audio: sound to play with video tensor
    """
    if is_notebook():
        from IPython.display import Video, display

        b = self.to_bytes(audio_tensor=audio)
        display(Video(data=b, embed=True, mimetype='video/mp4'))
    else:
        warnings.warn('Display of video is only possible in a notebook.')

@staticmethod
@abc.abstractmethod
def get_comp_backend() -> AbstractComputationalBackend:
    """The computational backend compatible with this tensor type."""
    ...

def save(
    self: 'T',
    file_path: Union[str, BytesIO],
    audio_tensor: Optional[AudioTensor] = None,
    video_frame_rate: int = 24,
    video_codec: str = 'h264',
    audio_frame_rate: int = 48000,
    audio_codec: str = 'aac',
    audio_format: str = 'fltp',
) -> None:
    """
    Save video tensor to a .mp4 file.

    ---

    ```python
    import numpy as np

    from docarray import BaseDoc
    from docarray.typing.tensor.audio.audio_tensor import AudioTensor
    from docarray.typing.tensor.video.video_tensor import VideoTensor


    class MyDoc(BaseDoc):
        video_tensor: VideoTensor
        audio_tensor: AudioTensor


    doc = MyDoc(
        video_tensor=np.random.randint(low=0, high=256, size=(10, 200, 300, 3)),
        audio_tensor=np.random.randn(100, 1, 1024).astype("float32"),
    )

    doc.video_tensor.save(
        file_path="/tmp/mp_.mp4",
        audio_tensor=doc.audio_tensor,
        audio_format="flt",
    )
    ```

    ---
    :param file_path: path to a .mp4 file. If file is a string, open the file by
        that name, otherwise treat it as a file-like object.
    :param audio_tensor: AudioTensor containing the video's soundtrack.
    :param video_frame_rate: video frames per second.
    :param video_codec: the name of a video decoder/encoder.
    :param audio_frame_rate: audio frames per second.
    :param audio_codec: the name of an audio decoder/encoder.
    :param audio_format: the name of one of the audio formats supported by PyAV,
        such as 'flt', 'fltp', 's16' or 's16p'.
    """
    if TYPE_CHECKING:
        import av
    else:
        av = import_library('av', raise_error=True)

    np_tensor = self.get_comp_backend().to_numpy(array=self)
    video_tensor = np_tensor.astype('uint8')

    if isinstance(file_path, str):
        format = file_path.split('.')[-1]
    else:
        format = 'mp4'

    with av.open(file_path, mode='w', format=format) as container:
        if video_tensor.ndim == 3:
            video_tensor = np.expand_dims(video_tensor, axis=0)

        stream_video = container.add_stream(video_codec, rate=video_frame_rate)
        stream_video.height = video_tensor.shape[-3]
        stream_video.width = video_tensor.shape[-2]

        if audio_tensor is not None:
            stream_audio = container.add_stream(audio_codec)
            audio_np = audio_tensor.get_comp_backend().to_numpy(array=audio_tensor)
            audio_layout = 'stereo' if audio_np.shape[-2] == 2 else 'mono'

            for i, audio in enumerate(audio_np):
                frame = av.AudioFrame.from_ndarray(
                    array=audio, format=audio_format, layout=audio_layout
                )
                frame.rate = audio_frame_rate
                frame.pts = audio.shape[-1] * i
                for packet in stream_audio.encode(frame):
                    container.mux(packet)

            for packet in stream_audio.encode(None):
                container.mux(packet)

        for vid in video_tensor:
            frame = av.VideoFrame.from_ndarray(vid, format='rgb24')
            for packet in stream_video.encode(frame):
                container.mux(packet)

        for packet in stream_video.encode(None):
            container.mux(packet)