import unittest

import numpy as np

from csst_proto.api import flip_image, read_test_image
from csst_proto.flip_image import flip_multiple_images_jl, flip_multiple_images_mp


class FlipImageTestCase(unittest.TestCase):
    def test_flip_image(self):
        """ test flip image """
        self.assertTrue(
            np.all(flip_image(read_test_image()) == np.array([[4, 3], [2, 1]]))
        )

        # the code fails for 1D array
        with self.assertRaises(AssertionError):
            flip_image(np.array([1, 2, 3, 4]))

    def test_flip_multiple_images_mp(self):
        """ test flip multiple images with multiprocessing """
        n_jobs = 10
        imgs = [read_test_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_mp(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(
                np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]]))
            )

    def test_flip_multiple_images_jl(self):
        """ test flip multiple images with joblib """
        n_jobs = 10
        imgs = [read_test_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_jl(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(
                np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]]))
            )