# Copyright 2019 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.
# ==============================================================================
"""
This is the test module for mindrecord
"""
import numpy as np
import os

import mindspore.dataset as ds
from mindspore import log as logger

from mindspore.mindrecord import FileWriter

FILES_NUM = 1
CV_FILE_NAME = "./complex.mindrecord"

def test_cv_minddataset_reader_multi_image_and_ndarray_tutorial():
    writer = FileWriter(CV_FILE_NAME, FILES_NUM)
    cv_schema_json={"id": {"type": "int32"},
                    "image_0": {"type": "bytes"},
                    "image_2": {"type": "bytes"},
                    "image_3": {"type": "bytes"},
                    "image_4": {"type": "bytes"},
                    "input_mask": {"type": "int32", "shape": [-1]},
                    "segments": {"type": "float32", "shape": [2, 3]}}
    writer.add_schema(cv_schema_json, "two_images_schema")
    with open("../data/mindrecord/testImageNetData/images/image_00010.jpg", "rb") as file_reader:
        img_data = file_reader.read()
    ndarray_1 = np.array([1, 2, 3, 4, 5], np.int32)
    ndarray_2 = np.array(([2, 3, 1], [7, 9, 0]), np.float32)
    data = []
    for i in range(5):
        item = {"id": i, "image_0": img_data, "image_2": img_data, "image_3": img_data, "image_4": img_data,
                "input_mask": ndarray_1, "segments": ndarray_2}
        data.append(item)
    writer.write_raw_data(data)
    writer.commit()
    assert os.path.exists(CV_FILE_NAME)
    assert os.path.exists(CV_FILE_NAME + ".db")

    """tutorial for minderdataset."""
    columns_list = ["id", "image_0", "image_2", "image_3", "image_4", "input_mask", "segments"]
    num_readers = 1
    data_set = ds.MindDataset(CV_FILE_NAME, columns_list, num_readers)
    assert data_set.get_dataset_size() == 5
    num_iter = 0
    for item in data_set.create_dict_iterator():
        assert len(item) == 7
        logger.info("item: {}".format(item))
        assert item["image_0"].dtype == np.uint8
        assert (item["image_0"] == item["image_2"]).all()
        assert (item["image_3"] == item["image_4"]).all()
        assert (item["image_0"] == item["image_4"]).all()
        assert item["image_2"].dtype == np.uint8
        assert item["image_3"].dtype == np.uint8
        assert item["image_4"].dtype == np.uint8
        assert item["id"].dtype == np.int32
        assert item["input_mask"].shape == (5, )
        assert item["input_mask"].dtype == np.int32
        assert item["segments"].shape == (2, 3)
        assert item["segments"].dtype == np.float32
        num_iter += 1
    assert num_iter == 5

    if os.path.exists("{}".format(CV_FILE_NAME + ".db")):
        os.remove(CV_FILE_NAME + ".db")
    if os.path.exists("{}".format(CV_FILE_NAME)):
        os.remove(CV_FILE_NAME)
