from __future__ import annotations

import unittest

import numpy as np
import torch

from qprefer_reward.processing import frame_to_pil, uniform_sample_frames


class ProcessingTest(unittest.TestCase):
    def test_minus_one_one_tensor_to_rgb(self) -> None:
        frame = torch.tensor(
            [
                [[-1.0, 1.0], [-1.0, 1.0]],
                [[-1.0, 1.0], [-1.0, 1.0]],
                [[-1.0, 1.0], [-1.0, 1.0]],
            ],
            dtype=torch.float32,
        )
        image = frame_to_pil(frame, value_range="minus_one_one")
        pixels = np.asarray(image)
        self.assertEqual(image.mode, "RGB")
        self.assertEqual(pixels.shape, (2, 2, 3))
        self.assertEqual(int(pixels.min()), 0)
        self.assertEqual(int(pixels.max()), 255)

    def test_uniform_sampler_repeats_short_video(self) -> None:
        video = np.zeros((2, 4, 4, 3), dtype=np.uint8)
        video[1] = 255
        frames = uniform_sample_frames(video, num_frames=8)
        self.assertEqual(len(frames), 8)
        self.assertEqual(int(np.asarray(frames[0]).mean()), 0)
        self.assertEqual(int(np.asarray(frames[-1]).mean()), 255)


if __name__ == "__main__":
    unittest.main()