# standard libraries
import contextlib
import datetime
import json
import logging
import os
import unittest
import uuid

# third party libraries
import numpy

# local libraries
from nion.data import Calibration
from nion.data import DataAndMetadata
from nion.swift.model import DataItem
from nion.swift.model import ImportExportManager
from nion.swift.model import Utility
from nion.swift.test import TestContext


class TestImportExportManagerClass(unittest.TestCase):

    def setUp(self):
        TestContext.begin_leaks()

    def tearDown(self):
        TestContext.end_leaks(self)

    def test_convert_data_element_records_time_zone_in_data_item_metadata(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element["datetime_modified"] = Utility.get_current_datetime_item()
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertIsNotNone(data_item.created)
            self.assertEqual(data_item.timezone_offset, data_element["datetime_modified"]["tz"])

    def test_convert_data_element_sets_timezone_and_timezone_offset_if_present(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element["datetime_modified"] = {'tz': '+0300', 'dst': '+60', 'local_datetime': '2015-06-10T19:31:52.780511', 'timezone': 'Europe/Athens'}
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertIsNotNone(data_item.created)
            self.assertEqual(data_item.timezone, "Europe/Athens")
            self.assertEqual(data_item.timezone_offset, "+0300")

    def test_date_formats(self):
        data_item = DataItem.DataItem()
        with contextlib.closing(data_item):
            data_item.created = datetime.datetime(2013, 11, 18, 14, 5, 4, 0)
            self.assertIsNotNone(data_item.created_local_as_string)

    def test_sub_area_size_change(self):
        with TestContext.create_memory_context() as test_context:
            document_model = test_context.create_document_model()
            data_element = dict()
            data_element["version"] = 1
            data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
            data_item = ImportExportManager.create_data_item_from_data_element(data_element)
            document_model.append_data_item(data_item)
            self.assertEqual(data_item.dimensional_shape, (16, 16))
            self.assertEqual(data_item.data_dtype, numpy.double)
            data_element["data"] = numpy.zeros((8, 8), dtype=numpy.double)
            data_element["sub_area"] = ((0,0), (4, 8))
            ImportExportManager.update_data_item_from_data_element(data_item, data_element)
            self.assertEqual(data_item.dimensional_shape, (8, 8))
            self.assertEqual(data_item.data_dtype, numpy.double)
            data_element["data"] = numpy.zeros((8, 8), dtype=float)
            data_element["sub_area"] = ((0,0), (4, 8))
            ImportExportManager.update_data_item_from_data_element(data_item, data_element)
            self.assertEqual(data_item.dimensional_shape, (8, 8))
            self.assertEqual(data_item.data_dtype, float)

    def test_ndata_write_to_then_read_from_temp_file(self):
        with TestContext.create_memory_context() as test_context:
            document_model = test_context.create_document_model()
            current_working_directory = os.getcwd()
            file_path = os.path.join(current_working_directory, "__file.ndata")
            handler = ImportExportManager.NDataImportExportHandler("ndata1-io-handler", "ndata", ["ndata"])
            data_item = DataItem.DataItem(numpy.zeros((16, 16), dtype=numpy.double))
            document_model.append_data_item(data_item)
            display_item = document_model.get_display_item_for_data_item(data_item)
            handler.write_display_item(None, display_item, file_path, "ndata")
            self.assertTrue(os.path.exists(file_path))
            try:
                data_items = handler.read_data_items(None, "ndata", file_path)
                self.assertEqual(len(data_items), 1)
                data_item = data_items[0]
                for data_item in data_items:
                    data_item.close()
            finally:
                os.remove(file_path)

    def test_npy_write_to_then_read_from_temp_file(self):
        with TestContext.create_memory_context() as test_context:
            document_model = test_context.create_document_model()
            current_working_directory = os.getcwd()
            file_path_npy = os.path.join(current_working_directory, "__file.npy")
            file_path_json = os.path.join(current_working_directory, "__file.json")
            numpy.save(file_path_npy, numpy.zeros((16, 16)))
            with open(file_path_json, "w") as f:
                json.dump({"version": 1}, f)
            handler = ImportExportManager.NumPyImportExportHandler("numpy-io-handler", "npy", ["npy"])
            try:
                data_items = handler.read_data_items(None, "npy", file_path_npy)
                self.assertEqual(len(data_items), 1)
                data_item = data_items[0]
                for data_item in data_items:
                    data_item.close()
            finally:
                os.remove(file_path_npy)
                os.remove(file_path_json)

    def test_get_writers_for_empty_data_item_returns_valid_list(self):
        data_item = DataItem.DataItem()
        with contextlib.closing(data_item):
            writers = ImportExportManager.ImportExportManager().get_writers_for_data_item(data_item)
            self.assertEqual(len(writers), 0)

    def test_get_writers_for_float_2d_data_item_returns_valid_list(self):
        data_item = DataItem.DataItem(numpy.zeros((8, 8), float))
        with contextlib.closing(data_item):
            writers = ImportExportManager.ImportExportManager().get_writers_for_data_item(data_item)
            self.assertTrue(len(writers) > 0)

    def test_data_element_date_gets_set_as_data_item_created_date(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element["datetime_modified"] = {'tz': '+0300', 'dst': '+60', 'local_datetime': '2015-06-10T19:31:52.780511'}
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertIsNotNone(data_item.created)
            self.assertEqual(data_item.timezone_offset, "+0300")
            local_offset_seconds = int(round((datetime.datetime.now() - datetime.datetime.utcnow()).total_seconds()))
            # check both matches for DST
            match1 = datetime.datetime(year=2015, month=6, day=10, hour=19 - 3, minute=31, second=52, microsecond=780511) + datetime.timedelta(seconds=local_offset_seconds)
            match2 = datetime.datetime(year=2015, month=6, day=10, hour=19 - 3, minute=31, second=52, microsecond=780511) + datetime.timedelta(seconds=local_offset_seconds + 3600)
            self.assertTrue(data_item.created_local == match1 or data_item.created_local == match2)

    def test_data_element_with_uuid_assigns_uuid_to_data_item(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element_uuid = uuid.uuid4()
        data_element["uuid"] = str(data_element_uuid)
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertEqual(data_item.uuid, data_element_uuid)

    def test_creating_data_element_with_sequence_data_makes_correct_data_item(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((4, 16, 16), dtype=numpy.double)
        data_element["is_sequence"] = True
        data_element["collection_dimension_count"] = 0
        data_element["datum_dimension_count"] = 2
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertEqual(data_item.is_sequence, True)
            self.assertEqual(data_item.collection_dimension_count, 0)
            self.assertEqual(data_item.datum_dimension_count, 2)
            self.assertEqual(data_item.xdata.is_sequence, True)
            self.assertEqual(data_item.xdata.collection_dimension_count, 0)
            self.assertEqual(data_item.xdata.datum_dimension_count, 2)

    def test_creating_data_element_with_sequence_and_implicit_datum_size_data_makes_correct_data_item(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((4, 16, 16), dtype=numpy.double)
        data_element["is_sequence"] = True
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertEqual(data_item.is_sequence, True)
            self.assertEqual(data_item.collection_dimension_count, 0)
            self.assertEqual(data_item.datum_dimension_count, 2)
            self.assertEqual(data_item.xdata.is_sequence, True)
            self.assertEqual(data_item.xdata.collection_dimension_count, 0)
            self.assertEqual(data_item.xdata.datum_dimension_count, 2)

    def test_data_element_to_extended_data_conversion(self):
        data = numpy.ones((8, 6), int)
        intensity_calibration = Calibration.Calibration(offset=1, scale=1.1, units="one")
        dimensional_calibrations = [Calibration.Calibration(offset=2, scale=2.1, units="two"), Calibration.Calibration(offset=3, scale=2.2, units="two")]
        metadata = {"hardware_source": {"one": 1, "two": "b"}}
        timestamp = datetime.datetime.now()
        data_descriptor = DataAndMetadata.DataDescriptor(is_sequence=False, collection_dimension_count=1, datum_dimension_count=1)
        xdata = DataAndMetadata.new_data_and_metadata(data, intensity_calibration=intensity_calibration, dimensional_calibrations=dimensional_calibrations, metadata=metadata, timestamp=timestamp, data_descriptor=data_descriptor)
        data_element = ImportExportManager.create_data_element_from_extended_data(xdata)
        new_xdata = ImportExportManager.convert_data_element_to_data_and_metadata(data_element)
        self.assertTrue(numpy.array_equal(data, new_xdata.data))
        self.assertNotEqual(id(new_xdata.intensity_calibration), id(intensity_calibration))
        self.assertEqual(new_xdata.intensity_calibration, intensity_calibration)
        self.assertNotEqual(id(new_xdata.dimensional_calibrations[0]), id(dimensional_calibrations[0]))
        self.assertEqual(new_xdata.dimensional_calibrations, dimensional_calibrations)
        self.assertNotEqual(id(new_xdata.metadata), id(metadata))
        self.assertEqual(new_xdata.metadata, metadata)
        self.assertNotEqual(id(new_xdata.data_descriptor), id(data_descriptor))
        self.assertEqual(new_xdata.data_descriptor, data_descriptor)

    def test_data_item_to_data_element_includes_time_zone(self):
        # created/modified are utc; timezone is specified in metadata/description/time_zone
        data_item = DataItem.DataItem(numpy.zeros((16, 16)))
        with contextlib.closing(data_item):
            data_item.created = datetime.datetime(2013, 6, 18, 14, 5, 4, 0)  # always utc
            data_item.timezone = "Europe/Athens"
            data_item.timezone_offset = "+0300"
            data_item._set_modified(datetime.datetime(2013, 6, 18, 14, 5, 4, 0))  # always utc
            data_item.metadata = {"description": {"time_zone": {"tz": "+0300", "dst": "+60"}}}
            data_element = ImportExportManager.create_data_element_from_data_item(data_item, include_data=False)
            self.assertEqual(data_element["datetime_modified"], {"dst": "+60", "local_datetime": "2013-06-18T17:05:04", 'tz': "+0300", 'timezone': "Europe/Athens"})

    def test_extended_data_to_data_element_includes_time_zone(self):
        # extended data timestamp is utc; timezone is specified in metadata/description/time_zone
        data = numpy.ones((8, 6), int)
        metadata = {"description": {"time_zone": {"tz": "+0300", "dst": "+60"}}}
        timestamp = datetime.datetime(2013, 11, 18, 14, 5, 4, 0)
        xdata = DataAndMetadata.new_data_and_metadata(data, metadata=metadata, timestamp=timestamp)
        data_element = ImportExportManager.create_data_element_from_extended_data(xdata)
        self.assertEqual(data_element["datetime_modified"], {"dst": "+60", "local_datetime": "2013-11-18T17:05:04", 'tz': "+0300"})

    def test_data_element_to_data_item_includes_time_zone(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element["datetime_modified"] = {'tz': '+0300', 'dst': '+60', 'local_datetime': '2015-06-10T19:31:52.780511'}
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertEqual(data_item.timezone_offset, "+0300")
            self.assertEqual(str(data_item.created), "2015-06-10 16:31:52.780511")

    def test_data_element_to_extended_data_includes_time_zone(self):
        data_element = dict()
        data_element["version"] = 1
        data_element["data"] = numpy.zeros((16, 16), dtype=numpy.double)
        data_element["datetime_modified"] = {'tz': '+0300', 'dst': '+60', 'local_datetime': '2015-06-10T19:31:52.780511'}
        xdata = ImportExportManager.convert_data_element_to_data_and_metadata(data_element)
        self.assertEqual(xdata.timezone_offset, "+0300")
        self.assertEqual(str(xdata.timestamp), "2015-06-10 16:31:52.780511")

    def test_time_zone_in_extended_data_to_data_element_to_data_item_conversion(self):
        # test the whole path, redundant?
        data = numpy.ones((8, 6), int)
        metadata = {"description": {"time_zone": {"tz": "+0300", "dst": "+60"}}, "hardware_source": {"one": 1, "two": "b"}}
        timestamp = datetime.datetime(2013, 11, 18, 14, 5, 4, 1)
        xdata = DataAndMetadata.new_data_and_metadata(data, metadata=metadata, timestamp=timestamp)
        data_element = ImportExportManager.create_data_element_from_extended_data(xdata)
        with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
            self.assertEqual(data_item.metadata["description"]["time_zone"]["tz"], "+0300")
            self.assertEqual(data_item.metadata["description"]["time_zone"]["dst"], "+60")
            self.assertEqual("2013-11-18 14:05:04.000001", str(data_item.created))

    def test_csv1_exporter_handles_multi_layer_display_item_with_same_calibration(self):
        with TestContext.create_memory_context() as test_context:
            document_model = test_context.create_document_model()
            data_item = DataItem.DataItem(numpy.full((50, ), 0, dtype=numpy.uint32))
            data_item.intensity_calibration = Calibration.Calibration(offset=10, scale=2)
            data_item2 = DataItem.DataItem(numpy.full((100, ), 1, dtype=numpy.uint32))
            data_item2.intensity_calibration = Calibration.Calibration(offset=5, scale=2)
            document_model.append_data_item(data_item)
            document_model.append_data_item(data_item2, False)
            display_item = document_model.get_display_item_for_data_item(data_item)
            display_item.display_type = "line_plot"
            display_item.append_display_data_channel_for_data_item(data_item2)
            current_working_directory = os.getcwd()
            file_path = os.path.join(current_working_directory, "__file.csv")
            handler = ImportExportManager.CSV1ImportExportHandler("csv1-io-handler", "CSV 1D", ["csv"])
            handler.write_display_item(None, display_item, file_path, "csv")
            self.assertTrue(os.path.exists(file_path))
            try:
                saved_data = numpy.genfromtxt(file_path, delimiter=", ")
                self.assertSequenceEqual(saved_data.shape, (max(data_item.xdata.data_shape[0], data_item2.xdata.data_shape[0]), 3))
                self.assertTrue(numpy.allclose(saved_data[:, 0], numpy.linspace(0, 99, 100)))
                self.assertTrue(numpy.allclose(saved_data[:50, 1], 10))
                self.assertTrue(numpy.allclose(saved_data[:, 2], 7))
            finally:
                os.remove(file_path)

    def test_csv1_exporter_handles_multi_layer_display_item_with_different_calibration(self):
        with TestContext.create_memory_context() as test_context:
            document_model = test_context.create_document_model()
            data_item = DataItem.DataItem(numpy.full((50, ), 0, dtype=numpy.uint32))
            data_item.dimensional_calibrations = [Calibration.Calibration(offset=0, scale=1, units="eV")]
            data_item.intensity_calibration = Calibration.Calibration(offset=10, scale=2)
            data_item2 = DataItem.DataItem(numpy.full((100, ), 1, dtype=numpy.uint32))
            data_item2.dimensional_calibrations = [Calibration.Calibration(offset=-10, scale=2, units="eV")]
            data_item2.intensity_calibration = Calibration.Calibration(offset=5, scale=2)
            document_model.append_data_item(data_item)
            document_model.append_data_item(data_item2, False)
            display_item = document_model.get_display_item_for_data_item(data_item)
            display_item.display_type = "line_plot"
            display_item.append_display_data_channel_for_data_item(data_item2)
            current_working_directory = os.getcwd()
            file_path = os.path.join(current_working_directory, "__file.csv")
            handler = ImportExportManager.CSV1ImportExportHandler("csv1-io-handler", "CSV 1D", ["csv"])
            handler.write_display_item(None, display_item, file_path, "csv")
            self.assertTrue(os.path.exists(file_path))
            try:
                saved_data = numpy.genfromtxt(file_path, delimiter=", ")
                self.assertSequenceEqual(saved_data.shape, (max(data_item.xdata.data_shape[0], data_item2.xdata.data_shape[0]), 4))
                self.assertTrue(numpy.allclose(saved_data[:50, 0], numpy.linspace(0, 49, 50)))
                self.assertTrue(numpy.allclose(saved_data[:50, 1], 10))
                self.assertTrue(numpy.allclose(saved_data[:, 2], numpy.linspace(-10, 188, 100)))
                self.assertTrue(numpy.allclose(saved_data[:, 3], 7))
            finally:
                os.remove(file_path)

    def test_data_item_to_data_element_produces_json_compatible_dict(self):
        data_item = DataItem.DataItem(numpy.zeros((16, 16)))
        with contextlib.closing(data_item):
            data_item.created = datetime.datetime(2013, 6, 18, 14, 5, 4, 0)  # always utc
            data_item.timezone = "Europe/Athens"
            data_item.timezone_offset = "+0300"
            data_item.source_file_path = "/path/to/source/file"
            data_item._set_modified(datetime.datetime(2013, 6, 18, 14, 5, 4, 0))  # always utc
            data_item.metadata = {"description": {"time_zone": {"tz": "+0300", "dst": "+60"}}}
            data_element = ImportExportManager.create_data_element_from_data_item(data_item, include_data=False)
            json.dumps(data_element)

    def test_data_item_to_data_element_and_back_keeps_large_format_flag(self):
        data_item = DataItem.DataItem(numpy.zeros((4, 4, 4)), large_format=True)
        with contextlib.closing(data_item):
            data_element = ImportExportManager.create_data_element_from_data_item(data_item, include_data=True)
            self.assertTrue(data_element.get("large_format"))
            with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
                self.assertTrue(data_item.large_format)

    def test_importing_rgb_does_not_set_large_format(self):
        data_item = DataItem.DataItem(numpy.zeros((8, 8, 4), dtype=float))
        with contextlib.closing(data_item):
            data_item_rgb = DataItem.DataItem(numpy.zeros((8, 8, 4), dtype=numpy.uint8))
            with contextlib.closing(data_item_rgb):
                data_element = ImportExportManager.create_data_element_from_data_item(data_item, include_data=True)
                data_element_rgb = ImportExportManager.create_data_element_from_data_item(data_item_rgb, include_data=True)
                data_element.pop("large_format")
                data_element_rgb.pop("large_format")
                with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element)) as data_item:
                    with contextlib.closing(ImportExportManager.create_data_item_from_data_element(data_element_rgb)) as data_item_rgb:
                        self.assertTrue(data_item.large_format)
                        self.assertFalse(data_item_rgb.large_format)

    def test_importing_large_numpy_file_sets_large_format_flag(self):
        current_working_directory = os.getcwd()
        file_path_npy = os.path.join(current_working_directory, "__file.npy")
        numpy.save(file_path_npy, numpy.zeros((4, 4, 4)))
        handler = ImportExportManager.NumPyImportExportHandler("numpy-io-handler", "npy", ["npy"])
        try:
            data_items = handler.read_data_items(None, "npy", file_path_npy)
            self.assertEqual(len(data_items), 1)
            data_item = data_items[0]
            self.assertTrue(data_item.large_format)
            for data_item in data_items:
                data_item.close()
        finally:
            os.remove(file_path_npy)

    def test_data_item_with_numpy_bool_to_data_element_produces_json_compatible_dict(self):
        data_item = DataItem.DataItem(numpy.zeros((16, 16)))
        data_item.large_format = numpy.product((2,3,4), dtype=numpy.int64) > 10  # produces a numpy.bool_
        with contextlib.closing(data_item):
            data_element = ImportExportManager.create_data_element_from_data_item(data_item, include_data=False)
            json.dumps(data_element)


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.DEBUG)
    unittest.main()
