import numpy as np
import torch
from PIL import Image

from stimulus_synthesis.media.normalize import videos_to_b_t_c_h_w


def test_videos_to_b_t_c_h_w_from_pil_frames():
    frames = [
        Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)),
        Image.fromarray(np.full((8, 8, 3), 255, dtype=np.uint8)),
    ]

    video = videos_to_b_t_c_h_w([frames], size=4, num_frames=3)

    assert video.shape == (1, 3, 3, 4, 4)
    assert torch.all(video >= 0)
    assert torch.all(video <= 1)


def test_videos_to_b_t_c_h_w_from_tensor_thwc():
    tensor = torch.zeros(2, 8, 8, 3)

    video = videos_to_b_t_c_h_w([tensor], size=4)

    assert video.shape == (1, 2, 3, 4, 4)


def test_videos_to_b_t_c_h_w_from_numpy_bthwc():
    array = np.zeros((1, 2, 8, 8, 3), dtype=np.float32)

    video = videos_to_b_t_c_h_w([array], size=4)

    assert video.shape == (1, 2, 3, 4, 4)