import unittest
import platform

from numpy import (
    uint8,
    uint16,
    uint32,
    uint64,
    zeros,
    float32,
    float64,
    all as alltrue,
    rint,
    arange,
)

import pygame
from pygame.locals import *

import pygame.surfarray


IS_PYPY = "PyPy" == platform.python_implementation()


@unittest.skipIf(IS_PYPY, "pypy skip known failure")  # TODO
class SurfarrayModuleTest(unittest.TestCase):
    pixels2d = {8: True, 16: True, 24: False, 32: True}
    pixels3d = {8: False, 16: False, 24: True, 32: True}
    array2d = {8: True, 16: True, 24: True, 32: True}
    array3d = {8: False, 16: False, 24: True, 32: True}

    test_palette = [
        (0, 0, 0, 255),
        (10, 30, 60, 255),
        (25, 75, 100, 255),
        (100, 150, 200, 255),
        (0, 100, 200, 255),
    ]
    surf_size = (10, 12)
    test_points = [
        ((0, 0), 1),
        ((4, 5), 1),
        ((9, 0), 2),
        ((5, 5), 2),
        ((0, 11), 3),
        ((4, 6), 3),
        ((9, 11), 4),
        ((5, 6), 4),
    ]

    @classmethod
    def setUpClass(cls):
        # Needed for 8 bits-per-pixel color palette surface tests.
        pygame.init()

    @classmethod
    def tearDownClass(cls):
        pygame.quit()

    def setUp(cls):
        # This makes sure pygame is always initialized before each test (in
        # case a test calls pygame.quit()).
        if not pygame.get_init():
            pygame.init()

    def _make_surface(self, bitsize, srcalpha=False, palette=None):
        if palette is None:
            palette = self.test_palette
        flags = 0
        if srcalpha:
            flags |= SRCALPHA
        surf = pygame.Surface(self.surf_size, flags, bitsize)
        if bitsize == 8:
            surf.set_palette([c[:3] for c in palette])
        return surf

    def _fill_surface(self, surf, palette=None):
        if palette is None:
            palette = self.test_palette
        surf.fill(palette[1], (0, 0, 5, 6))
        surf.fill(palette[2], (5, 0, 5, 6))
        surf.fill(palette[3], (0, 6, 5, 6))
        surf.fill(palette[4], (5, 6, 5, 6))

    def _make_src_surface(self, bitsize, srcalpha=False, palette=None):
        surf = self._make_surface(bitsize, srcalpha, palette)
        self._fill_surface(surf, palette)
        return surf

    def _assert_surface(self, surf, palette=None, msg=""):
        if palette is None:
            palette = self.test_palette
        if surf.get_bitsize() == 16:
            palette = [surf.unmap_rgb(surf.map_rgb(c)) for c in palette]
        for posn, i in self.test_points:
            self.assertEqual(
                surf.get_at(posn),
                palette[i],
                "%s != %s: flags: %i, bpp: %i, posn: %s%s"
                % (
                    surf.get_at(posn),
                    palette[i],
                    surf.get_flags(),
                    surf.get_bitsize(),
                    posn,
                    msg,
                ),
            )

    def _make_array3d(self, dtype):
        return zeros((self.surf_size[0], self.surf_size[1], 3), dtype)

    def _fill_array2d(self, arr, surf):
        palette = self.test_palette
        arr[:5, :6] = surf.map_rgb(palette[1]) & 0xFFFFFFFF
        arr[5:, :6] = surf.map_rgb(palette[2]) & 0xFFFFFFFF
        arr[:5, 6:] = surf.map_rgb(palette[3]) & 0xFFFFFFFF
        arr[5:, 6:] = surf.map_rgb(palette[4]) & 0xFFFFFFFF

    def _fill_array3d(self, arr):
        palette = self.test_palette
        arr[:5, :6] = palette[1][:3]
        arr[5:, :6] = palette[2][:3]
        arr[:5, 6:] = palette[3][:3]
        arr[5:, 6:] = palette[4][:3]

    def _make_src_array3d(self, dtype):
        arr = self._make_array3d(dtype)
        self._fill_array3d(arr)
        return arr

    def _make_array2d(self, dtype):
        return zeros(self.surf_size, dtype)

    def test_array2d(self):
        sources = [
            self._make_src_surface(8),
            self._make_src_surface(16),
            self._make_src_surface(16, srcalpha=True),
            self._make_src_surface(24),
            self._make_src_surface(32),
            self._make_src_surface(32, srcalpha=True),
        ]
        palette = self.test_palette
        alpha_color = (0, 0, 0, 128)

        for surf in sources:
            arr = pygame.surfarray.array2d(surf)
            for posn, i in self.test_points:
                self.assertEqual(
                    arr[posn],
                    surf.get_at_mapped(posn),
                    "%s != %s: flags: %i, bpp: %i, posn: %s"
                    % (
                        arr[posn],
                        surf.get_at_mapped(posn),
                        surf.get_flags(),
                        surf.get_bitsize(),
                        posn,
                    ),
                )

            if surf.get_masks()[3]:
                surf.fill(alpha_color)
                arr = pygame.surfarray.array2d(surf)
                posn = (0, 0)
                self.assertEqual(
                    arr[posn],
                    surf.get_at_mapped(posn),
                    "%s != %s: bpp: %i"
                    % (arr[posn], surf.get_at_mapped(posn), surf.get_bitsize()),
                )

    def test_array3d(self):
        sources = [
            self._make_src_surface(16),
            self._make_src_surface(16, srcalpha=True),
            self._make_src_surface(24),
            self._make_src_surface(32),
            self._make_src_surface(32, srcalpha=True),
        ]
        palette = self.test_palette

        for surf in sources:
            arr = pygame.surfarray.array3d(surf)

            def same_color(ac, sc):
                return ac[0] == sc[0] and ac[1] == sc[1] and ac[2] == sc[2]

            for posn, i in self.test_points:
                self.assertTrue(
                    same_color(arr[posn], surf.get_at(posn)),
                    "%s != %s: flags: %i, bpp: %i, posn: %s"
                    % (
                        tuple(arr[posn]),
                        surf.get_at(posn),
                        surf.get_flags(),
                        surf.get_bitsize(),
                        posn,
                    ),
                )

    def test_array_alpha(self):
        palette = [
            (0, 0, 0, 0),
            (10, 50, 100, 255),
            (60, 120, 240, 130),
            (64, 128, 255, 0),
            (255, 128, 0, 65),
        ]
        targets = [
            self._make_src_surface(8, palette=palette),
            self._make_src_surface(16, palette=palette),
            self._make_src_surface(16, palette=palette, srcalpha=True),
            self._make_src_surface(24, palette=palette),
            self._make_src_surface(32, palette=palette),
            self._make_src_surface(32, palette=palette, srcalpha=True),
        ]

        for surf in targets:
            p = palette
            if surf.get_bitsize() == 16:
                p = [surf.unmap_rgb(surf.map_rgb(c)) for c in p]
            arr = pygame.surfarray.array_alpha(surf)
            if surf.get_masks()[3]:
                for (x, y), i in self.test_points:
                    self.assertEqual(
                        arr[x, y],
                        p[i][3],
                        (
                            "%i != %i, posn: (%i, %i), "
                            "bitsize: %i"
                            % (arr[x, y], p[i][3], x, y, surf.get_bitsize())
                        ),
                    )
            else:
                self.assertTrue(alltrue(arr == 255))

        # No per-pixel alpha when blanket alpha is None.
        for surf in targets:
            blanket_alpha = surf.get_alpha()
            surf.set_alpha(None)
            arr = pygame.surfarray.array_alpha(surf)
            self.assertTrue(
                alltrue(arr == 255),
                "All alpha values should be 255 when"
                " surf.set_alpha(None) has been set."
                " bitsize: %i, flags: %i" % (surf.get_bitsize(), surf.get_flags()),
            )
            surf.set_alpha(blanket_alpha)

        # Bug for per-pixel alpha surface when blanket alpha 0.
        for surf in targets:
            blanket_alpha = surf.get_alpha()
            surf.set_alpha(0)
            arr = pygame.surfarray.array_alpha(surf)
            if surf.get_masks()[3]:
                self.assertFalse(
                    alltrue(arr == 255),
                    "bitsize: %i, flags: %i" % (surf.get_bitsize(), surf.get_flags()),
                )
            else:
                self.assertTrue(
                    alltrue(arr == 255),
                    "bitsize: %i, flags: %i" % (surf.get_bitsize(), surf.get_flags()),
                )
            surf.set_alpha(blanket_alpha)

    def test_array_colorkey(self):
        palette = [
            (0, 0, 0, 0),
            (10, 50, 100, 255),
            (60, 120, 240, 130),
            (64, 128, 255, 0),
            (255, 128, 0, 65),
        ]
        targets = [
            self._make_src_surface(8, palette=palette),
            self._make_src_surface(16, palette=palette),
            self._make_src_surface(16, palette=palette, srcalpha=True),
            self._make_src_surface(24, palette=palette),
            self._make_src_surface(32, palette=palette),
            self._make_src_surface(32, palette=palette, srcalpha=True),
        ]

        for surf in targets:
            p = palette
            if surf.get_bitsize() == 16:
                p = [surf.unmap_rgb(surf.map_rgb(c)) for c in p]
            surf.set_colorkey(None)
            arr = pygame.surfarray.array_colorkey(surf)
            self.assertTrue(alltrue(arr == 255))

            for i in range(1, len(palette)):
                surf.set_colorkey(p[i])
                alphas = [255] * len(p)
                alphas[i] = 0
                arr = pygame.surfarray.array_colorkey(surf)
                for (x, y), j in self.test_points:
                    self.assertEqual(
                        arr[x, y],
                        alphas[j],
                        (
                            "%i != %i, posn: (%i, %i), "
                            "bitsize: %i"
                            % (arr[x, y], alphas[j], x, y, surf.get_bitsize())
                        ),
                    )

    def test_array_red(self):
        self._test_array_rgb("red", 0)

    def test_array_green(self):
        self._test_array_rgb("green", 1)

    def test_array_blue(self):
        self._test_array_rgb("blue", 2)

    def _test_array_rgb(self, operation, mask_posn):
        method_name = "array_" + operation

        array_rgb = getattr(pygame.surfarray, method_name)
        palette = [
            (0, 0, 0, 255),
            (5, 13, 23, 255),
            (29, 31, 37, 255),
            (131, 157, 167, 255),
            (179, 191, 251, 255),
        ]
        plane = [c[mask_posn] for c in palette]

        targets = [
            self._make_src_surface(24, palette=palette),
            self._make_src_surface(32, palette=palette),
            self._make_src_surface(32, palette=palette, srcalpha=True),
        ]

        for surf in targets:
            self.assertFalse(surf.get_locked())
            for (x, y), i in self.test_points:
                surf.fill(palette[i])
                arr = array_rgb(surf)
                self.assertEqual(arr[x, y], plane[i])
                surf.fill((100, 100, 100, 250))
                self.assertEqual(arr[x, y], plane[i])
                self.assertFalse(surf.get_locked())
                del arr

    def test_blit_array(self):
        s = pygame.Surface((10, 10), 0, 24)
        a = pygame.surfarray.array3d(s)
        pygame.surfarray.blit_array(s, a)

        # target surfaces
        targets = [
            self._make_surface(8),
            self._make_surface(16),
            self._make_surface(16, srcalpha=True),
            self._make_surface(24),
            self._make_surface(32),
            self._make_surface(32, srcalpha=True),
        ]

        # source arrays
        arrays3d = []
        dtypes = [(8, uint8), (16, uint16), (32, uint32)]
        try:
            dtypes.append((64, uint64))
        except NameError:
            pass
        arrays3d = [(self._make_src_array3d(dtype), None) for __, dtype in dtypes]
        for bitsize in [8, 16, 24, 32]:
            palette = None
            if bitsize == 16:
                s = pygame.Surface((1, 1), 0, 16)
                palette = [s.unmap_rgb(s.map_rgb(c)) for c in self.test_palette]
            if self.pixels3d[bitsize]:
                surf = self._make_src_surface(bitsize)
                arr = pygame.surfarray.pixels3d(surf)
                arrays3d.append((arr, palette))
            if self.array3d[bitsize]:
                surf = self._make_src_surface(bitsize)
                arr = pygame.surfarray.array3d(surf)
                arrays3d.append((arr, palette))
                for sz, dtype in dtypes:
                    arrays3d.append((arr.astype(dtype), palette))

        # tests on arrays
        def do_blit(surf, arr):
            pygame.surfarray.blit_array(surf, arr)

        for surf in targets:
            bitsize = surf.get_bitsize()
            for arr, palette in arrays3d:
                surf.fill((0, 0, 0, 0))
                if bitsize == 8:
                    self.assertRaises(ValueError, do_blit, surf, arr)
                else:
                    pygame.surfarray.blit_array(surf, arr)
                    self._assert_surface(surf, palette)

            if self.pixels2d[bitsize]:
                surf.fill((0, 0, 0, 0))
                s = self._make_src_surface(bitsize, surf.get_flags() & SRCALPHA)
                arr = pygame.surfarray.pixels2d(s)
                pygame.surfarray.blit_array(surf, arr)
                self._assert_surface(surf)

            if self.array2d[bitsize]:
                s = self._make_src_surface(bitsize, surf.get_flags() & SRCALPHA)
                arr = pygame.surfarray.array2d(s)
                for sz, dtype in dtypes:
                    surf.fill((0, 0, 0, 0))
                    if sz >= bitsize:
                        pygame.surfarray.blit_array(surf, arr.astype(dtype))
                        self._assert_surface(surf)
                    else:
                        self.assertRaises(
                            ValueError, do_blit, surf, self._make_array2d(dtype)
                        )

        # Check alpha for 2D arrays
        surf = self._make_surface(16, srcalpha=True)
        arr = zeros(surf.get_size(), uint16)
        arr[...] = surf.map_rgb((0, 128, 255, 64))
        color = surf.unmap_rgb(arr[0, 0])
        pygame.surfarray.blit_array(surf, arr)
        self.assertEqual(surf.get_at((5, 5)), color)

        surf = self._make_surface(32, srcalpha=True)
        arr = zeros(surf.get_size(), uint32)
        color = (0, 111, 255, 63)
        arr[...] = surf.map_rgb(color)
        pygame.surfarray.blit_array(surf, arr)
        self.assertEqual(surf.get_at((5, 5)), color)

        # Check shifts
        arr3d = self._make_src_array3d(uint8)

        shift_tests = [
            (16, [12, 0, 8, 4], [0xF000, 0xF, 0xF00, 0xF0]),
            (24, [16, 0, 8, 0], [0xFF0000, 0xFF, 0xFF00, 0]),
            (32, [0, 16, 24, 8], [0xFF, 0xFF0000, 0xFF000000, 0xFF00]),
        ]

        for bitsize, shifts, masks in shift_tests:
            surf = self._make_surface(bitsize, srcalpha=(shifts[3] != 0))
            palette = None
            if bitsize == 16:
                palette = [surf.unmap_rgb(surf.map_rgb(c)) for c in self.test_palette]

            self.assertRaises(TypeError, surf.set_shifts, shifts)
            self.assertRaises(TypeError, surf.set_masks, masks)

        # Invalid arrays
        surf = pygame.Surface((1, 1), 0, 32)
        t = "abcd"
        self.assertRaises(ValueError, do_blit, surf, t)

        surf_size = self.surf_size
        surf = pygame.Surface(surf_size, 0, 32)
        arr = zeros([surf_size[0], surf_size[1] + 1, 3], uint32)
        self.assertRaises(ValueError, do_blit, surf, arr)
        arr = zeros([surf_size[0] + 1, surf_size[1], 3], uint32)
        self.assertRaises(ValueError, do_blit, surf, arr)

        surf = pygame.Surface((1, 4), 0, 32)
        arr = zeros((4,), uint32)
        self.assertRaises(ValueError, do_blit, surf, arr)
        arr.shape = (1, 1, 1, 4)
        self.assertRaises(ValueError, do_blit, surf, arr)

        # Issue #81: round from float to int
        try:
            rint
        except NameError:
            pass
        else:
            surf = pygame.Surface((10, 10), pygame.SRCALPHA, 32)
            w, h = surf.get_size()
            length = w * h
            for dtype in [float32, float64]:
                surf.fill((255, 255, 255, 0))
                farr = arange(0, length, dtype=dtype)
                farr.shape = w, h
                pygame.surfarray.blit_array(surf, farr)
                for x in range(w):
                    for y in range(h):
                        self.assertEqual(
                            surf.get_at_mapped((x, y)), int(rint(farr[x, y]))
                        )

    # this test should be removed soon, when the function is deleted
    def test_get_arraytype(self):
        array_type = pygame.surfarray.get_arraytype()

        self.assertEqual(array_type, "numpy", f"unknown array type {array_type}")

    # this test should be removed soon, when the function is deleted
    def test_get_arraytypes(self):
        arraytypes = pygame.surfarray.get_arraytypes()
        self.assertIn("numpy", arraytypes)

        for atype in arraytypes:
            self.assertEqual(atype, "numpy", f"unknown array type {atype}")

    def test_make_surface(self):
        # How does one properly test this with 2d arrays. It makes no sense
        # since the pixel format is not entirely dependent on element size.
        # Just make sure the surface pixel size is at least as large as the
        # array element size I guess.
        #
        for bitsize, dtype in [(8, uint8), (16, uint16), (24, uint32)]:
            ## Even this simple assertion fails for 2d arrays. Where's the problem?
            ##            surf = pygame.surfarray.make_surface(self._make_array2d(dtype))
            ##            self.assertGreaterEqual(surf.get_bitsize(), bitsize,
            ##                            "not %i >= %i)" % (surf.get_bitsize(), bitsize))
            ##
            surf = pygame.surfarray.make_surface(self._make_src_array3d(dtype))
            self._assert_surface(surf)

        # Issue #81: round from float to int
        try:
            rint
        except NameError:
            pass
        else:
            w = 9
            h = 11
            length = w * h
            for dtype in [float32, float64]:
                farr = arange(0, length, dtype=dtype)
                farr.shape = w, h
                surf = pygame.surfarray.make_surface(farr)
                for x in range(w):
                    for y in range(h):
                        self.assertEqual(
                            surf.get_at_mapped((x, y)), int(rint(farr[x, y]))
                        )

    def test_map_array(self):
        arr3d = self._make_src_array3d(uint8)
        targets = [
            self._make_surface(8),
            self._make_surface(16),
            self._make_surface(16, srcalpha=True),
            self._make_surface(24),
            self._make_surface(32),
            self._make_surface(32, srcalpha=True),
        ]
        palette = self.test_palette

        for surf in targets:
            arr2d = pygame.surfarray.map_array(surf, arr3d)
            for posn, i in self.test_points:
                self.assertEqual(
                    arr2d[posn],
                    surf.map_rgb(palette[i]),
                    "%i != %i, bitsize: %i, flags: %i"
                    % (
                        arr2d[posn],
                        surf.map_rgb(palette[i]),
                        surf.get_bitsize(),
                        surf.get_flags(),
                    ),
                )

        # Exception checks
        self.assertRaises(
            ValueError,
            pygame.surfarray.map_array,
            self._make_surface(32),
            self._make_array2d(uint8),
        )

    def test_pixels2d(self):
        sources = [
            self._make_surface(8),
            self._make_surface(16, srcalpha=True),
            self._make_surface(32, srcalpha=True),
        ]

        for surf in sources:
            self.assertFalse(surf.get_locked())
            arr = pygame.surfarray.pixels2d(surf)
            self.assertTrue(surf.get_locked())
            self._fill_array2d(arr, surf)
            surf.unlock()
            self.assertTrue(surf.get_locked())
            del arr
            self.assertFalse(surf.get_locked())
            self.assertEqual(surf.get_locks(), ())
            self._assert_surface(surf)

        # Error checks
        self.assertRaises(ValueError, pygame.surfarray.pixels2d, self._make_surface(24))

    def test_pixels3d(self):
        sources = [self._make_surface(24), self._make_surface(32)]

        for surf in sources:
            self.assertFalse(surf.get_locked())
            arr = pygame.surfarray.pixels3d(surf)
            self.assertTrue(surf.get_locked())
            self._fill_array3d(arr)
            surf.unlock()
            self.assertTrue(surf.get_locked())
            del arr
            self.assertFalse(surf.get_locked())
            self.assertEqual(surf.get_locks(), ())
            self._assert_surface(surf)

        # Alpha check
        color = (1, 2, 3, 0)
        surf = self._make_surface(32, srcalpha=True)
        arr = pygame.surfarray.pixels3d(surf)
        arr[0, 0] = color[:3]
        self.assertEqual(surf.get_at((0, 0)), color)

        # Error checks
        def do_pixels3d(surf):
            pygame.surfarray.pixels3d(surf)

        self.assertRaises(ValueError, do_pixels3d, self._make_surface(8))
        self.assertRaises(ValueError, do_pixels3d, self._make_surface(16))

    def test_pixels_alpha(self):
        palette = [
            (0, 0, 0, 0),
            (127, 127, 127, 0),
            (127, 127, 127, 85),
            (127, 127, 127, 170),
            (127, 127, 127, 255),
        ]
        alphas = [0, 45, 86, 99, 180]

        surf = self._make_src_surface(32, srcalpha=True, palette=palette)

        self.assertFalse(surf.get_locked())
        arr = pygame.surfarray.pixels_alpha(surf)
        self.assertTrue(surf.get_locked())
        surf.unlock()
        self.assertTrue(surf.get_locked())

        for (x, y), i in self.test_points:
            self.assertEqual(arr[x, y], palette[i][3])

        for (x, y), i in self.test_points:
            alpha = alphas[i]
            arr[x, y] = alpha
            color = (127, 127, 127, alpha)
            self.assertEqual(surf.get_at((x, y)), color, "posn: (%i, %i)" % (x, y))

        del arr
        self.assertFalse(surf.get_locked())
        self.assertEqual(surf.get_locks(), ())

        # Check exceptions.
        def do_pixels_alpha(surf):
            pygame.surfarray.pixels_alpha(surf)

        targets = [(8, False), (16, False), (16, True), (24, False), (32, False)]

        for bitsize, srcalpha in targets:
            self.assertRaises(
                ValueError, do_pixels_alpha, self._make_surface(bitsize, srcalpha)
            )

    def test_pixels_red(self):
        self._test_pixels_rgb("red", 0)

    def test_pixels_green(self):
        self._test_pixels_rgb("green", 1)

    def test_pixels_blue(self):
        self._test_pixels_rgb("blue", 2)

    def _test_pixels_rgb(self, operation, mask_posn):
        method_name = "pixels_" + operation

        pixels_rgb = getattr(pygame.surfarray, method_name)
        palette = [
            (0, 0, 0, 255),
            (5, 13, 23, 255),
            (29, 31, 37, 255),
            (131, 157, 167, 255),
            (179, 191, 251, 255),
        ]
        plane = [c[mask_posn] for c in palette]

        surf24 = self._make_src_surface(24, srcalpha=False, palette=palette)
        surf32 = self._make_src_surface(32, srcalpha=False, palette=palette)
        surf32a = self._make_src_surface(32, srcalpha=True, palette=palette)

        for surf in [surf24, surf32, surf32a]:
            self.assertFalse(surf.get_locked())
            arr = pixels_rgb(surf)
            self.assertTrue(surf.get_locked())
            surf.unlock()
            self.assertTrue(surf.get_locked())

            for (x, y), i in self.test_points:
                self.assertEqual(arr[x, y], plane[i])

            del arr
            self.assertFalse(surf.get_locked())
            self.assertEqual(surf.get_locks(), ())

        # Check exceptions.
        targets = [(8, False), (16, False), (16, True)]

        for bitsize, srcalpha in targets:
            self.assertRaises(
                ValueError, pixels_rgb, self._make_surface(bitsize, srcalpha)
            )

    def test_use_arraytype(self):
        def do_use_arraytype(atype):
            pygame.surfarray.use_arraytype(atype)

        pygame.surfarray.use_arraytype("numpy")
        self.assertEqual(pygame.surfarray.get_arraytype(), "numpy")
        self.assertRaises(ValueError, do_use_arraytype, "not an option")

    def test_surf_lock(self):
        sf = pygame.Surface((5, 5), 0, 32)
        for atype in pygame.surfarray.get_arraytypes():
            pygame.surfarray.use_arraytype(atype)

            ar = pygame.surfarray.pixels2d(sf)
            self.assertTrue(sf.get_locked())

            sf.unlock()
            self.assertTrue(sf.get_locked())

            del ar
            self.assertFalse(sf.get_locked())
            self.assertEqual(sf.get_locks(), ())


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