"""Test yuv_tools helper functions."""

import os
import pathlib
from typing import Tuple
import unittest

from avm_stats import yuv_tools
import numpy as np

TESTDATA_ENV_VAR = "LIBAOM_TEST_DATA_PATH"
TEST_STREAM_WIDTH = 160
TEST_STREAM_HEIGHT = 90


class YuvToolsTest(unittest.TestCase):

  def setUp(self):
    super().setUp()
    if TESTDATA_ENV_VAR not in os.environ:
      raise RuntimeError(
          f"{TESTDATA_ENV_VAR} environment variable must be set before running"
          " pytest!"
      )
    self.testdata_path = pathlib.Path(os.environ[TESTDATA_ENV_VAR])

  def test_upscale_plane(self):
    arr = np.array([[1, 2], [3, 4]])
    expected = np.array([
        [1, 1, 2, 2],
        [1, 1, 2, 2],
        [3, 3, 4, 4],
        [3, 3, 4, 4],
    ])
    upscaled = yuv_tools.upscale(arr, 2)
    self.assertTrue(np.array_equal(upscaled, expected))

  def test_yuv_to_rgb(self):
    y = np.array([[50, 100], [150, 200]])
    u = np.array([[128, 64], [128, 192]])
    v = np.array([[128, 128], [64, 128]])
    expected = np.array([
        [[50, 49, 51], [100, 124, 0]],
        [[77, 186, 151], [200, 174, 255]],
    ])
    rgb = yuv_tools.yuv_to_rgb(y, u, v)
    self.assertTrue(np.array_equal(rgb, expected))

  def check_yuv_seq(
      self,
      yuv_seq: yuv_tools.YuvSequence,
      expected_frames: int,
      expected_first_pixel_rgb: Tuple[int, int, int],
  ):
    self.assertEqual(len(yuv_seq.yuvs), expected_frames)
    for i in range(expected_frames):
      self.assertEqual(
          yuv_seq.yuvs[i].rgb.shape, (TEST_STREAM_HEIGHT, TEST_STREAM_WIDTH, 3)
      )

    first_pixel_rgb = tuple(yuv_seq.yuvs[0].rgb[0, 0])
    self.assertEqual(first_pixel_rgb, expected_first_pixel_rgb)

  def test_parse_y4m(self):
    num_frames = 3
    yuv_path = self.testdata_path / "park_joy_90p_8_420.y4m"
    yuv_seq = yuv_tools.parse_raw_yuv(
        yuv_path,
        width=TEST_STREAM_WIDTH,
        height=TEST_STREAM_HEIGHT,
        num_frames=3,
    )
    self.check_yuv_seq(
        yuv_seq, num_frames, expected_first_pixel_rgb=(70, 133, 0)
    )