import contextlib
import io
from pathlib import Path
import tempfile
import unittest

import cv2
import numpy as np

from colorchecker import ResultCode, detect_colorchecker, read_image


class ColorCheckerTests(unittest.TestCase):
    def assert_failure(self, path, expected):
        code, rgb, images = detect_colorchecker(path)
        self.assertEqual(code, expected)
        self.assertEqual(rgb.shape, (3, 24))
        self.assertEqual(rgb.dtype, np.float64)
        self.assertTrue(np.isnan(rgb).all())
        self.assertEqual(images, {})

    def assert_images(self, images, input_shape):
        self.assertEqual(set(images), {'warped', 'mask', 'debug'})
        self.assertEqual(images['warped'].shape, (560, 840, 3))
        self.assertEqual(images['debug'].shape, (560, 840, 3))
        self.assertEqual(images['mask'].shape, input_shape[:2])
        for image in images.values():
            self.assertEqual(image.dtype, np.uint8)
        self.assertTrue(np.isin(images['mask'], [0, 255]).all())
        self.assertFalse(np.array_equal(images['warped'], images['debug']))
        self.assertFalse(np.shares_memory(images['warped'], images['debug']))

    def test_missing_file(self):
        with tempfile.TemporaryDirectory() as directory:
            self.assert_failure(Path(directory) / 'missing.png', ResultCode.IMAGE_READ_ERROR)

    def test_invalid_and_empty_file(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'invalid.png'
            for data in (b'', b'not an image'):
                path.write_bytes(data)
                self.assert_failure(path, ResultCode.IMAGE_READ_ERROR)

    def test_no_checker(self):
        self.assert_failure(np.zeros((560, 840, 3), dtype=np.uint8),
                            ResultCode.DETECTION_FAILED)
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'blank.png'
            cv2.imwrite(str(path), np.zeros((560, 840, 3), dtype=np.uint8))
            self.assert_failure(path, ResultCode.DETECTION_FAILED)

    def test_invalid_image_arrays(self):
        for image in (
            np.zeros((0, 10, 3), dtype=np.uint8),
            np.zeros((10, 10), dtype=np.uint8),
            np.zeros((10, 10, 4), dtype=np.uint8),
            np.zeros((10, 10, 3), dtype=np.float32),
            np.array(1, dtype=np.uint8),
            None,
            123,
        ):
            with self.subTest(shape=getattr(image, 'shape', None)):
                self.assert_failure(image, ResultCode.IMAGE_READ_ERROR)

    def test_rgb_layout_and_unicode_path_without_output(self):
        colors = np.array([
            [115, 82, 68], [194, 150, 130], [98, 122, 157],
            [87, 108, 67], [133, 128, 177], [103, 189, 170],
            [214, 126, 44], [80, 91, 166], [193, 90, 99],
            [94, 60, 108], [157, 188, 64], [224, 163, 46],
            [56, 61, 150], [70, 148, 73], [175, 54, 60],
            [231, 199, 31], [187, 86, 149], [8, 133, 161],
            [243, 243, 243], [200, 200, 200], [160, 160, 160],
            [122, 122, 122], [85, 85, 85], [52, 52, 52],
        ], dtype=np.uint8)
        image = np.zeros((560, 840, 3), dtype=np.uint8)
        for index, rgb in enumerate(colors):
            row, col = divmod(index, 6)
            image[row * 140 + 25:row * 140 + 115,
                  col * 140 + 25:col * 140 + 115] = rgb[::-1]
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / '色票 画像.png'
            path.write_bytes(cv2.imencode('.png', image)[1].tobytes())
            stdout = io.StringIO()
            with contextlib.redirect_stdout(stdout):
                code, measured, images = detect_colorchecker(path)
            self.assertEqual(code, ResultCode.SUCCESS)
            self.assertEqual(measured.dtype, np.float64)
            self.assert_images(images, image.shape)
            np.testing.assert_allclose(measured, colors.T, atol=1)
            self.assertEqual(stdout.getvalue(), '')
            self.assertEqual(list(Path(directory).iterdir()), [path])
            # 明るい背景に囲まれた色票と、縦向き・反転を含む全方向。
            surrounded = np.pad(image, ((30, 30), (30, 30), (0, 0)),
                                constant_values=240)
            sources = [str(path), image, image[:, ::-1]]
            sources.extend(np.rot90(surrounded, k) for k in range(4))
            sources.extend(np.rot90(surrounded[:, ::-1], k) for k in range(4))
            for source in sources:
                with self.subTest(input_type=type(source)):
                    original = source.copy() if isinstance(source, np.ndarray) else None
                    code, measured, images = detect_colorchecker(filename=source)
                    self.assertEqual(code, ResultCode.SUCCESS)
                    self.assert_images(images, source.shape if original is not None
                                       else image.shape)
                    np.testing.assert_allclose(measured, colors.T, atol=1)
                    if original is not None:
                        np.testing.assert_array_equal(source, original)

    def test_tias26_checker(self):
        path = Path(__file__).parent / 'sample' / 'TIAS26_MCC.jpg'
        image = read_image(path)
        code, measured, images = detect_colorchecker(image)
        self.assertEqual(code, ResultCode.SUCCESS)
        self.assert_images(images, image.shape)

        # 原画像で確認した24色票の内部と比較し、位置・向き・順序も検証。
        expected = np.empty((3, 24))
        for row, y in enumerate((600, 900, 1190, 1475, 1750, 2020)):
            for col, x in enumerate((480, 800, 1110, 1420)):
                rgb = image[y - 20:y + 20, x - 20:x + 20, ::-1].mean(axis=(0, 1))
                expected[:, (3 - col) * 6 + row] = rgb
        np.testing.assert_allclose(measured, expected, atol=8)


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