# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Testing Invert op in DE
"""
import numpy as np

import mindspore.dataset as ds
import mindspore.dataset.transforms.py_transforms
import mindspore.dataset.vision.py_transforms as F
import mindspore.dataset.vision.c_transforms as C
from mindspore import log as logger
from util import visualize_list, save_and_check_md5, diff_mse

DATA_DIR = "../data/dataset/testImageNetData/train/"

GENERATE_GOLDEN = False


def test_invert_py(plot=False):
    """
    Test Invert python op
    """
    logger.info("Test Invert Python op")

    # Original Images
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transforms_original = mindspore.dataset.transforms.py_transforms.Compose([F.Decode(),
                                                                              F.Resize((224, 224)),
                                                                              F.ToTensor()])

    ds_original = data_set.map(operations=transforms_original, input_columns="image")

    ds_original = ds_original.batch(512)

    for idx, (image, _) in enumerate(ds_original):
        if idx == 0:
            images_original = np.transpose(image.asnumpy(), (0, 2, 3, 1))
        else:
            images_original = np.append(images_original,
                                        np.transpose(image.asnumpy(), (0, 2, 3, 1)),
                                        axis=0)

    # Color Inverted Images
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transforms_invert = mindspore.dataset.transforms.py_transforms.Compose([F.Decode(),
                                                                            F.Resize((224, 224)),
                                                                            F.Invert(),
                                                                            F.ToTensor()])

    ds_invert = data_set.map(operations=transforms_invert, input_columns="image")

    ds_invert = ds_invert.batch(512)

    for idx, (image, _) in enumerate(ds_invert):
        if idx == 0:
            images_invert = np.transpose(image.asnumpy(), (0, 2, 3, 1))
        else:
            images_invert = np.append(images_invert,
                                      np.transpose(image.asnumpy(), (0, 2, 3, 1)),
                                      axis=0)

    num_samples = images_original.shape[0]
    mse = np.zeros(num_samples)
    for i in range(num_samples):
        mse[i] = np.mean((images_invert[i] - images_original[i]) ** 2)
    logger.info("MSE= {}".format(str(np.mean(mse))))

    if plot:
        visualize_list(images_original, images_invert)


def test_invert_c(plot=False):
    """
    Test Invert Cpp op
    """
    logger.info("Test Invert cpp op")

    # Original Images
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transforms_original = [C.Decode(), C.Resize(size=[224, 224])]

    ds_original = data_set.map(operations=transforms_original, input_columns="image")

    ds_original = ds_original.batch(512)

    for idx, (image, _) in enumerate(ds_original):
        if idx == 0:
            images_original = image.asnumpy()
        else:
            images_original = np.append(images_original,
                                        image.asnumpy(),
                                        axis=0)

    # Invert Images
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transform_invert = [C.Decode(), C.Resize(size=[224, 224]),
                        C.Invert()]

    ds_invert = data_set.map(operations=transform_invert, input_columns="image")

    ds_invert = ds_invert.batch(512)

    for idx, (image, _) in enumerate(ds_invert):
        if idx == 0:
            images_invert = image.asnumpy()
        else:
            images_invert = np.append(images_invert,
                                      image.asnumpy(),
                                      axis=0)
    if plot:
        visualize_list(images_original, images_invert)

    num_samples = images_original.shape[0]
    mse = np.zeros(num_samples)
    for i in range(num_samples):
        mse[i] = diff_mse(images_invert[i], images_original[i])
    logger.info("MSE= {}".format(str(np.mean(mse))))


def test_invert_py_c(plot=False):
    """
    Test Invert Cpp op and python op
    """
    logger.info("Test Invert cpp and python op")

    # Invert Images in cpp
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)
    data_set = data_set.map(operations=[C.Decode(), C.Resize((224, 224))], input_columns=["image"])

    ds_c_invert = data_set.map(operations=C.Invert(), input_columns="image")

    ds_c_invert = ds_c_invert.batch(512)

    for idx, (image, _) in enumerate(ds_c_invert):
        if idx == 0:
            images_c_invert = image.asnumpy()
        else:
            images_c_invert = np.append(images_c_invert,
                                        image.asnumpy(),
                                        axis=0)

    # invert images in python
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)
    data_set = data_set.map(operations=[C.Decode(), C.Resize((224, 224))], input_columns=["image"])

    transforms_p_invert = mindspore.dataset.transforms.py_transforms.Compose([lambda img: img.astype(np.uint8),
                                                                              F.ToPIL(),
                                                                              F.Invert(),
                                                                              np.array])

    ds_p_invert = data_set.map(operations=transforms_p_invert, input_columns="image")

    ds_p_invert = ds_p_invert.batch(512)

    for idx, (image, _) in enumerate(ds_p_invert):
        if idx == 0:
            images_p_invert = image.asnumpy()
        else:
            images_p_invert = np.append(images_p_invert,
                                        image.asnumpy(),
                                        axis=0)

    num_samples = images_c_invert.shape[0]
    mse = np.zeros(num_samples)
    for i in range(num_samples):
        mse[i] = diff_mse(images_p_invert[i], images_c_invert[i])
    logger.info("MSE= {}".format(str(np.mean(mse))))

    if plot:
        visualize_list(images_c_invert, images_p_invert, visualize_mode=2)


def test_invert_one_channel():
    """
     Test Invert cpp op with one channel image
     """
    logger.info("Test Invert C Op With One Channel Images")

    c_op = C.Invert()

    try:
        data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)
        data_set = data_set.map(operations=[C.Decode(), C.Resize((224, 224)),
                                            lambda img: np.array(img[:, :, 0])], input_columns=["image"])

        data_set.map(operations=c_op, input_columns="image")

    except RuntimeError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "The shape" in str(e)


def test_invert_md5_py():
    """
    Test Invert python op with md5 check
    """
    logger.info("Test Invert python op with md5 check")

    # Generate dataset
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transforms_invert = mindspore.dataset.transforms.py_transforms.Compose([F.Decode(),
                                                                            F.Invert(),
                                                                            F.ToTensor()])

    data = data_set.map(operations=transforms_invert, input_columns="image")
    # Compare with expected md5 from images
    filename = "invert_01_result_py.npz"
    save_and_check_md5(data, filename, generate_golden=GENERATE_GOLDEN)


def test_invert_md5_c():
    """
    Test Invert cpp op with md5 check
    """
    logger.info("Test Invert cpp op with md5 check")

    # Generate dataset
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)

    transforms_invert = [C.Decode(),
                         C.Resize(size=[224, 224]),
                         C.Invert(),
                         F.ToTensor()]

    data = data_set.map(operations=transforms_invert, input_columns="image")
    # Compare with expected md5 from images
    filename = "invert_01_result_c.npz"
    save_and_check_md5(data, filename, generate_golden=GENERATE_GOLDEN)


if __name__ == "__main__":
    test_invert_py(plot=False)
    test_invert_c(plot=False)
    test_invert_py_c(plot=False)
    test_invert_one_channel()
    test_invert_md5_py()
    test_invert_md5_c()

[((36, 4, 36, 40), 'mindspore.log.info', 'logger.info', ({(36, 16, 36, 39): '"""Test Invert Python op"""'}, {}), "('Test Invert Python op')", True, 'from mindspore import log as logger\n'), ((39, 15, 39, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((58, 15, 58, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((78, 10, 78, 31), 'numpy.zeros', 'np.zeros', ({(78, 19, 78, 30): 'num_samples'}, {}), '(num_samples)', True, 'import numpy as np\n'), ((91, 4, 91, 37), 'mindspore.log.info', 'logger.info', ({(91, 16, 91, 36): '"""Test Invert cpp op"""'}, {}), "('Test Invert cpp op')", True, 'from mindspore import log as logger\n'), ((94, 15, 94, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((111, 15, 111, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((131, 10, 131, 31), 'numpy.zeros', 'np.zeros', ({(131, 19, 131, 30): 'num_samples'}, {}), '(num_samples)', True, 'import numpy as np\n'), ((141, 4, 141, 48), 'mindspore.log.info', 'logger.info', ({(141, 16, 141, 47): '"""Test Invert cpp and python op"""'}, {}), "('Test Invert cpp and python op')", True, 'from mindspore import log as logger\n'), ((144, 15, 144, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((160, 15, 160, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((181, 10, 181, 31), 'numpy.zeros', 'np.zeros', ({(181, 19, 181, 30): 'num_samples'}, {}), '(num_samples)', True, 'import numpy as np\n'), ((194, 4, 194, 59), 'mindspore.log.info', 'logger.info', ({(194, 16, 194, 58): '"""Test Invert C Op With One Channel Images"""'}, {}), "('Test Invert C Op With One Channel Images')", True, 'from mindspore import log as logger\n'), ((196, 11, 196, 21), 'mindspore.dataset.vision.c_transforms.Invert', 'C.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((214, 4, 214, 55), 'mindspore.log.info', 'logger.info', ({(214, 16, 214, 54): '"""Test Invert python op with md5 check"""'}, {}), "('Test Invert python op with md5 check')", True, 'from mindspore import log as logger\n'), ((217, 15, 217, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((226, 4, 226, 71), 'util.save_and_check_md5', 'save_and_check_md5', (), '', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((233, 4, 233, 52), 'mindspore.log.info', 'logger.info', ({(233, 16, 233, 51): '"""Test Invert cpp op with md5 check"""'}, {}), "('Test Invert cpp op with md5 check')", True, 'from mindspore import log as logger\n'), ((236, 15, 236, 73), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((246, 4, 246, 71), 'util.save_and_check_md5', 'save_and_check_md5', (), '', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((80, 17, 80, 70), 'numpy.mean', 'np.mean', ({(80, 25, 80, 69): '(images_invert[i] - images_original[i]) ** 2'}, {}), '((images_invert[i] - images_original[i]) ** 2)', True, 'import numpy as np\n'), ((84, 8, 84, 54), 'util.visualize_list', 'visualize_list', ({(84, 23, 84, 38): 'images_original', (84, 40, 84, 53): 'images_invert'}, {}), '(images_original, images_invert)', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((96, 27, 96, 37), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((96, 39, 96, 64), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', (), '', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((113, 24, 113, 34), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((113, 36, 113, 61), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', (), '', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((114, 24, 114, 34), 'mindspore.dataset.vision.c_transforms.Invert', 'C.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((128, 8, 128, 54), 'util.visualize_list', 'visualize_list', ({(128, 23, 128, 38): 'images_original', (128, 40, 128, 53): 'images_invert'}, {}), '(images_original, images_invert)', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((133, 17, 133, 63), 'util.diff_mse', 'diff_mse', ({(133, 26, 133, 42): 'images_invert[i]', (133, 44, 133, 62): 'images_original[i]'}, {}), '(images_invert[i], images_original[i])', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((183, 17, 183, 65), 'util.diff_mse', 'diff_mse', ({(183, 26, 183, 44): 'images_p_invert[i]', (183, 46, 183, 64): 'images_c_invert[i]'}, {}), '(images_p_invert[i], images_c_invert[i])', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((187, 8, 187, 74), 'util.visualize_list', 'visualize_list', (), '', False, 'from util import visualize_list, save_and_check_md5, diff_mse\n'), ((199, 19, 199, 77), 'mindspore.dataset.ImageFolderDataset', 'ds.ImageFolderDataset', (), '', True, 'import mindspore.dataset as ds\n'), ((238, 25, 238, 35), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((239, 25, 239, 50), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', (), '', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((240, 25, 240, 35), 'mindspore.dataset.vision.c_transforms.Invert', 'C.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((241, 25, 241, 37), 'mindspore.dataset.vision.py_transforms.ToTensor', 'F.ToTensor', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((41, 78, 41, 88), 'mindspore.dataset.vision.py_transforms.Decode', 'F.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((42, 78, 42, 98), 'mindspore.dataset.vision.py_transforms.Resize', 'F.Resize', ({(42, 87, 42, 97): '(224, 224)'}, {}), '((224, 224))', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((43, 78, 43, 90), 'mindspore.dataset.vision.py_transforms.ToTensor', 'F.ToTensor', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((60, 76, 60, 86), 'mindspore.dataset.vision.py_transforms.Decode', 'F.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((61, 76, 61, 96), 'mindspore.dataset.vision.py_transforms.Resize', 'F.Resize', ({(61, 85, 61, 95): '(224, 224)'}, {}), '((224, 224))', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((62, 76, 62, 86), 'mindspore.dataset.vision.py_transforms.Invert', 'F.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((63, 76, 63, 88), 'mindspore.dataset.vision.py_transforms.ToTensor', 'F.ToTensor', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((147, 42, 147, 52), 'mindspore.dataset.vision.c_transforms.Invert', 'C.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((164, 78, 164, 87), 'mindspore.dataset.vision.py_transforms.ToPIL', 'F.ToPIL', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((165, 78, 165, 88), 'mindspore.dataset.vision.py_transforms.Invert', 'F.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((219, 76, 219, 86), 'mindspore.dataset.vision.py_transforms.Decode', 'F.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((220, 76, 220, 86), 'mindspore.dataset.vision.py_transforms.Invert', 'F.Invert', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((221, 76, 221, 88), 'mindspore.dataset.vision.py_transforms.ToTensor', 'F.ToTensor', ({}, {}), '()', True, 'import mindspore.dataset.vision.py_transforms as F\n'), ((81, 37, 81, 49), 'numpy.mean', 'np.mean', ({(81, 45, 81, 48): 'mse'}, {}), '(mse)', True, 'import numpy as np\n'), ((134, 37, 134, 49), 'numpy.mean', 'np.mean', ({(134, 45, 134, 48): 'mse'}, {}), '(mse)', True, 'import numpy as np\n'), ((145, 40, 145, 50), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((145, 52, 145, 72), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', ({(145, 61, 145, 71): '(224, 224)'}, {}), '((224, 224))', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((161, 40, 161, 50), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((161, 52, 161, 72), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', ({(161, 61, 161, 71): '(224, 224)'}, {}), '((224, 224))', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((184, 37, 184, 49), 'numpy.mean', 'np.mean', ({(184, 45, 184, 48): 'mse'}, {}), '(mse)', True, 'import numpy as np\n'), ((200, 44, 200, 54), 'mindspore.dataset.vision.c_transforms.Decode', 'C.Decode', ({}, {}), '()', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((200, 56, 200, 76), 'mindspore.dataset.vision.c_transforms.Resize', 'C.Resize', ({(200, 65, 200, 75): '(224, 224)'}, {}), '((224, 224))', True, 'import mindspore.dataset.vision.c_transforms as C\n'), ((201, 56, 201, 78), 'numpy.array', 'np.array', ({(201, 65, 201, 77): 'img[:, :, (0)]'}, {}), '(img[:, :, (0)])', True, 'import numpy as np\n')]