import numpy as np
from pathlib import Path
import tempfile
import warnings

from pynwb import NWBHDF5IO, validate, TimeSeries
from pynwb.device import DeviceModel
from pynwb.ecephys import ElectrodesTable
from pynwb.image import ImageSeries
from pynwb.misc import FrequencyBandsTable
from pynwb.testing import TestCase


class TestReadOldVersions(TestCase):

    expected_warnings = {
        ''
        '2.1.0_imageseries_non_external_format.nwb': [(
            "ImageSeries 'test_imageseries': Format must be 'external' when external_file is specified."
        )],
        '2.1.0_imageseries_nonmatch_starting_frame.nwb': [(
            "ImageSeries 'test_imageseries': The number of frame indices in 'starting_frame' should have the same "
            "length as 'external_file'."
        )],
    }

    expected_errors = {
        '1.0.2_str_experimenter.nwb': [("root/general/experimenter (general/experimenter): incorrect shape - expected "
                                        "an array of shape '[None]', got non-array data 'one experimenter'")],
        '1.0.3_str_experimenter.nwb': [("root/general/experimenter (general/experimenter): incorrect shape - expected "
                                        "an array of shape '[None]', got non-array data 'one experimenter'")],
        '1.0.2_str_pub.nwb': [("root/general/related_publications (general/related_publications): incorrect shape "
                               "- expected an array of shape '[None]', got non-array data 'one publication'")],
        '1.0.3_str_pub.nwb': [("root/general/related_publications (general/related_publications): incorrect shape "
                               "- expected an array of shape '[None]', got non-array data 'one publication'")],
        '1.5.1_timeseries_no_data.nwb': [("TimeSeries/data/data (acquisition/test_timeseries/data): argument missing")],
        '1.5.1_timeseries_no_unit.nwb': [("TimeSeries/data/unit (acquisition/test_timeseries/data): argument missing")],
        '1.5.1_imageseries_no_data.nwb': [("ImageSeries/data/data (acquisition/test_imageseries/data): "
                                           "argument missing")],
        '1.5.1_imageseries_no_unit.nwb': [("ImageSeries/data/unit (acquisition/test_imageseries/data): "
                                           "argument missing")],
    }

    def get_io(self, path):
        """Get an NWBHDF5IO object for the given path."""
        return NWBHDF5IO(str(path), 'r')

    def test_read(self):
        """Test reading and validating all NWB files in the same folder as this file.

        This folder contains NWB files generated by previous versions of NWB using the script
        src/pynwb/testing/make_test_files.py
        """
        dir_path = Path(__file__).parent
        nwb_files = dir_path.glob('*.nwb')
        for f in nwb_files:
            with self.subTest(file=f.name):
                with warnings.catch_warnings(record=True) as warnings_on_read:
                    warnings.simplefilter("always")
                    with self.get_io(f) as io:
                        errors = validate(io)
                        io.read()
                        for w in warnings_on_read:
                            if f.name in self.expected_warnings:
                                if str(w.message) not in self.expected_warnings[f.name]:
                                    pass
                                    # will replace above with below after the test file is updated
                                    # raise Exception("Unexpected warning: %s: %s" % (f.name, str(w.message)))
                            else:
                                pass
                                # will replace above with below after the test file is updated
                                # raise Exception("Unexpected warning: %s: %s" % (f.name, str(w.message)))
                        if errors:
                            for e in errors:
                                if f.name in self.expected_errors:
                                    if str(e) not in self.expected_errors[f.name]:
                                        warnings.warn('%s: %s' % (f.name, e))
                                else:
                                    raise Exception("Unexpected validation error: %s: %s" % (f.name, e))
                            # TODO uncomment below when validation errors have been fixed
                            # raise Exception('%d validation error(s). See warnings.' % len(errors))

    def test_read_timeseries_no_data(self):
        """Test that a TimeSeries written without data is read with data set to the default value."""
        f = Path(__file__).parent / '1.5.1_timeseries_no_data.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            np.testing.assert_array_equal(read_nwbfile.acquisition['test_timeseries'].data, TimeSeries.DEFAULT_DATA)

    def test_read_timeseries_no_unit(self):
        """Test that an ImageSeries written without unit is read with unit set to the default value."""
        f = Path(__file__).parent / '1.5.1_timeseries_no_unit.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            self.assertEqual(read_nwbfile.acquisition['test_timeseries'].unit, TimeSeries.DEFAULT_UNIT)

    def test_read_imageseries_no_data(self):
        """Test that an ImageSeries written without data is read with data set to the default value."""
        f = Path(__file__).parent / '1.5.1_imageseries_no_data.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            np.testing.assert_array_equal(read_nwbfile.acquisition['test_imageseries'].data, ImageSeries.DEFAULT_DATA)

    def test_read_imageseries_no_unit(self):
        """Test that an ImageSeries written without unit is read with unit set to the default value."""
        f = Path(__file__).parent / '1.5.1_imageseries_no_unit.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            self.assertEqual(read_nwbfile.acquisition['test_imageseries'].unit, ImageSeries.DEFAULT_UNIT)

    def test_read_imageseries_non_external_format(self):
        """Test that reading an ImageSeries with an inconsistent format does not change the value."""
        fbase = "2.1.0_imageseries_non_external_format.nwb"
        f = Path(__file__).parent / fbase
        expected_warning = self.expected_warnings[fbase][0]
        with self.assertWarnsWith(UserWarning, expected_warning):
            with self.get_io(f) as io:
                read_nwbfile = io.read()
                self.assertEqual(read_nwbfile.acquisition['test_imageseries'].format, "tiff")

    def test_read_imageseries_nonmatch_starting_frame(self):
        """Test that reading an ImageSeries with an inconsistent starting_frame does not change the value."""
        fbase = "2.1.0_imageseries_nonmatch_starting_frame.nwb"
        f = Path(__file__).parent / fbase
        expected_warning = self.expected_warnings[fbase][0]
        with self.assertWarnsWith(UserWarning, expected_warning):
            with self.get_io(f) as io:
                read_nwbfile = io.read()
                np.testing.assert_array_equal(read_nwbfile.acquisition['test_imageseries'].starting_frame, [1, 2, 3])

    def test_read_subject_no_age__reference(self):
        """Test that reading a Subject without an age__reference set with NWB schema 2.5.0 sets the value to None"""
        f = Path(__file__).parent / '2.2.0_subject_no_age__reference.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            self.assertIsNone(read_nwbfile.subject.age__reference)

    def test_read_electrodes_table_as_neurodata_type(self):
        """Test that an "electrodes" table written as a DynamicTable is read as an ElectrodesTable"""
        f = Path(__file__).parent / '3.0.0_electrodes_dynamic_table.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            assert isinstance(read_nwbfile.electrodes, ElectrodesTable)

            # test that references to the electrodes table are also ElectrodesTable
            assert read_nwbfile.units.electrodes.table is read_nwbfile.electrodes
            assert read_nwbfile.acquisition["ElectricalSeries"].electrodes.table is read_nwbfile.electrodes

            # test that export writes the correct builders
            with tempfile.TemporaryDirectory() as temp_dir:
                export_file = Path(temp_dir) / "3.0.0_electrodes_dynamic_table_export.nwb"
                with NWBHDF5IO(export_file, 'w') as export_io:
                    export_io.export(io)

                with self.get_io(export_file) as read_export_io:
                    read_export_nwbfile = read_export_io.read()
                    assert isinstance(read_export_nwbfile.electrodes, ElectrodesTable)

                    # test that references to the electrodes table are also ElectrodesTable
                    units_table_ref = read_export_nwbfile.units.electrodes.table
                    assert units_table_ref is read_export_nwbfile.electrodes
                    eseries_table_ref = read_export_nwbfile.acquisition["ElectricalSeries"].electrodes.table
                    assert eseries_table_ref is read_export_nwbfile.electrodes

    def test_read_bands_table_as_neurodata_type(self):
        """Test that a "bands" table written as a DynamicTable is read as an FrequencyBandsTable"""
        f = Path(__file__).parent / '3.0.0_decompositionseries_bands_dynamic_table.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()
            assert isinstance(read_nwbfile.processing['test_mod']['LFPSpectralAnalysis'].bands, FrequencyBandsTable)
    
    def test_read_device_model_str_attribute(self):
        """Test that a Device.model written as a string attribute is read and remapped to a DeviceModel object"""
        f = Path(__file__).parent / '3.0.0_fiber_photometry_extension.nwb'
        with self.get_io(f) as io:
            # assert warning is issued to inform user the attribute is being remapped
            with self.assertWarnsWith(UserWarning,
                                      'Device.model was detected as a string, ' \
                                      'but NWB 2.9 specifies Device.model as a link to a DeviceModel. '
                                      'Remapping "dichroic mirror model" to a new DeviceModel.'):
                read_nwbfile = io.read()

            # assert data was remapped correctly
            device = read_nwbfile.devices['dichroic_mirror_1']
            self.assertIsInstance(device.model, DeviceModel)
            self.assertEqual(device.model.name, 'dichroic mirror model')
            self.assertEqual(device.model.description, 'Dichroic mirror for green indicator')
            self.assertEqual(device.model.manufacturer, '')
            self.assertEqual(device.model.model_number, None)
    
    def test_read_device_model_link_to_other_object(self):
        """Test that a Device.model written as a link to another object is read and remapped to a new attribute"""
        f = Path(__file__).parent / '3.0.0_optogenetics_extension.nwb'
        with self.get_io(f) as io:
            # assert warning is issued to inform user where old data is being remapped
            with self.assertWarnsWith(UserWarning,
                                      'The model attribute of the Device "Lambda" was detected as a non-DeviceModel '
                                      'object. Data associated with this object can be accessed at '
                                      '\"nwbfile.devices["Lambda"].ndx_optogenetics_model\"'):
                read_nwbfile = io.read()

            # assert data was remapped correctly
            device = read_nwbfile.devices['Lambda']
            self.assertIsNone(device.model)
            self.assertIsNotNone(device.ndx_optogenetics_model)
            self.assertEqual(device.ndx_optogenetics_model.name, 'Lambda Model')
            self.assertEqual(device.ndx_optogenetics_model.description, 'Lambda fiber (tapered fiber) from Optogenix.')
            self.assertEqual(device.ndx_optogenetics_model.numerical_aperture, 0.39)

    def test_read_device_model_link_to_extension_device_model(self):
        """Test that a Device.model written as a link to an extension DeviceModel object is read successfully"""
        f = Path(__file__).parent / '3.0.0_ophys_devices_extension.nwb'
        with self.get_io(f) as io:
            read_nwbfile = io.read()

            # assert model is read correct
            band_optical_filter = read_nwbfile.devices['band_optical_filter']
            self.assertIsInstance(band_optical_filter.model, DeviceModel)
            self.assertEqual(band_optical_filter.model.name, 'band_optical_filter_model')
            self.assertEqual(band_optical_filter.model.description, 'Band optical filter model for green indicator')
            self.assertEqual(band_optical_filter.model.bandwidth_in_nm, 30.0)
