"""Tests for data conversion related functionality."""
import os
import shutil
import tempfile
from typing import Tuple, Union
import unittest
import warnings

from pathlib import Path

from mockito import mock, when, unstub, verify, verifyNoMoreInteractions, any as any_value

from bcipy.config import DEFAULT_ENCODING, RAW_DATA_FILENAME, TRIGGER_FILENAME, DEFAULT_PARAMETER_FILENAME
from bcipy.helpers import convert
from bcipy.helpers.convert import (
    archive_list,
    compress,
    convert_to_bdf,
    convert_to_edf,
    convert_to_mne,
    decompress,
    pyedf_convert,
    tobii_to_norm,
    norm_to_tobii,
)
from bcipy.helpers.parameters import Parameters
from bcipy.helpers.raw_data import sample_data, write, RawData
from bcipy.helpers.triggers import MOCK_TRIGGER_DATA
from bcipy.signal.generator.generator import gen_random_data

from mne.io import read_raw_edf, read_raw_bdf


def create_bcipy_session_artifacts(
        write_dir: str,
        channels: Union[int, list] = 3,
        sample_rate: int = 300,
        samples: int = 5000,
        filter_settings: dict = {
            'filter_low': 0.5,
            'filter_high': 30,
            'filter_order': 5,
            'notch_filter_frequency': 60,
            'down_sampling_rate': 3,
            'static_trigger_offset': 0.0
        },
) -> Tuple[str, RawData, Parameters]:
    """Write BciPy session artifacts to a temporary directory.

    This includes a raw data file, trigger file, and a parameters file.
    """
    trg_data = MOCK_TRIGGER_DATA
    if isinstance(channels, int):
        channels = [f'ch{i}' for i in range(channels)]
    data = sample_data(ch_names=channels, sample_rate=sample_rate, rows=samples)
    with open(Path(write_dir, TRIGGER_FILENAME), 'w', encoding=DEFAULT_ENCODING) as trg_file:
        trg_file.write(trg_data)

    write(data, Path(write_dir, f'{RAW_DATA_FILENAME}.csv'))

    params = Parameters.from_cast_values(raw_data_name=f'{RAW_DATA_FILENAME}.csv',
                                         trigger_file_name=TRIGGER_FILENAME,
                                         preview_inquiry_length=0.5,
                                         time_fixation=0.5,
                                         time_prompt=0.5,
                                         time_flash=0.5,
                                         # define filter settings
                                         static_trigger_offset=filter_settings['static_trigger_offset'],
                                         down_sampling_rate=filter_settings['down_sampling_rate'],
                                         notch_filter_frequency=filter_settings['notch_filter_frequency'],
                                         filter_high=filter_settings['filter_high'],
                                         filter_low=filter_settings['filter_low'],
                                         filter_order=filter_settings['filter_order'])
    params.save(write_dir, DEFAULT_PARAMETER_FILENAME)
    return trg_data, data, params


class TestEDFConvert(unittest.TestCase):
    """Tests for EDF data format conversions."""

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp()
        _, data, _ = create_bcipy_session_artifacts(self.temp_dir)
        self.channels = data.channels

    def tearDown(self):
        shutil.rmtree(self.temp_dir)
        verifyNoMoreInteractions()
        unstub()

    def test_convert_to_edf_defaults(self):
        """Test default behavior"""
        path = convert_to_edf(self.temp_dir)
        self.assertTrue(os.path.exists(path))

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            edf = read_raw_edf(path, preload=True)

        self.assertTrue(len(edf.get_data()) > 0)

        for ch_name in self.channels:
            self.assertTrue(ch_name in edf.ch_names)

    def test_convert_to_edf_overwrite_false(self):
        """Test overwriting fails"""

        convert_to_edf(self.temp_dir)
        with self.assertRaises(OSError):
            convert_to_edf(self.temp_dir, overwrite=False)

    def test_convert_to_edf_overwrite_true(self):
        """Test that overwriting can be configured"""

        convert_to_edf(self.temp_dir)
        convert_to_edf(self.temp_dir, overwrite=True)

    def test_convert_to_edf_with_custom_path(self):
        """Test creating the EDF using a custom edf path"""
        path = convert_to_edf(self.temp_dir,
                              edf_path=str(Path(self.temp_dir, 'mydata.edf').resolve()))

        self.assertEqual(Path(path).name, 'mydata.edf')

    def test_convert_to_edf_raises_value_error_with_invalid_edfpath(self):
        """Test that an value error is raised when the edf path is invalid"""
        edf_path = 'invalid.pdf'
        with self.assertRaises(ValueError):
            convert_to_edf(self.temp_dir, edf_path=edf_path)

    def test_convert_to_edf_with_write_targetness(self):
        """Test creating the EDF using targetness for event annotations"""
        path = convert_to_edf(self.temp_dir,
                              write_targetness=True)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            edf = read_raw_edf(path, preload=True)

        self.assertIn('target', edf.annotations.description)
        self.assertIn('nontarget', edf.annotations.description)

    def test_convert_to_edf_without_write_targetness(self):
        """Test creating the EDF with labels as event annotations"""
        path = convert_to_edf(self.temp_dir,
                              write_targetness=False)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            edf = read_raw_edf(path, preload=True)

        self.assertNotIn('target', edf.annotations.description)
        self.assertNotIn('nontarget', edf.annotations.description)
        self.assertIn('+', edf.annotations.description)

    def test_convert_to_edf_with_pre_filter(self):
        """Test creating the EDF with pre-filtering"""

        data = mock()
        channels = mock()
        fs = 100
        events = mock()
        annotation_channels = 1
        return_path = mock()
        when(convert).pyedf_convert(
            any_value(str),
            write_targetness=any_value(bool),
            use_event_durations=any_value(bool),
            remove_pre_fixation=any_value(bool),
            pre_filter=True,
        ).thenReturn((
            data,
            channels,
            fs,
            events,
            annotation_channels,
            'filter_settings'
        ))
        when(convert).write_pyedf(
            any_value(str),
            data,
            channels,
            fs,
            events,
            any_value(int),
            any_value(bool),
            annotation_channels,
            any_value(str)).thenReturn(return_path)
        path = convert_to_edf(self.temp_dir,
                              pre_filter=True)

        verify(convert, times=1).pyedf_convert(
            any_value(str),
            write_targetness=any_value(bool),
            use_event_durations=any_value(bool),
            remove_pre_fixation=any_value(bool),
            pre_filter=True,
        )
        verify(convert, times=1).write_pyedf(
            any_value(str),
            data,
            channels,
            fs,
            events,
            any_value(int),
            any_value(bool),
            annotation_channels,
            any_value(str))
        self.assertEqual(path, return_path)


class TestBDFConvert(unittest.TestCase):
    """Tests for BDF data format conversions."""

    def setUp(self):
        """Override; set up the needed path for load functions."""

        self.temp_dir = tempfile.mkdtemp()
        _, data, _ = create_bcipy_session_artifacts(self.temp_dir)
        self.channels = data.channels

    def tearDown(self):
        """Override"""
        shutil.rmtree(self.temp_dir)
        verifyNoMoreInteractions()
        unstub()

    def test_convert_to_bdf_defaults(self):
        """Test default convert to bdf behavior"""
        path = convert_to_bdf(self.temp_dir)
        self.assertTrue(os.path.exists(path))

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            bdf = read_raw_bdf(path, preload=True)

        self.assertTrue(len(bdf.get_data()) > 0)

        for ch_name in self.channels:
            self.assertTrue(ch_name in bdf.ch_names)

    def test_convert_to_bdf_overwrite_false(self):
        """Test overwriting fails if not configured"""

        convert_to_bdf(self.temp_dir)
        with self.assertRaises(OSError):
            convert_to_bdf(self.temp_dir, overwrite=False)

    def test_convert_to_bdf_overwrite_true(self):
        """Test that overwriting can be configured"""

        convert_to_bdf(self.temp_dir)
        convert_to_bdf(self.temp_dir, overwrite=True)

    def test_convert_to_bdf_with_custom_path(self):
        """Test creating the EDF using a custom edf path"""
        path = convert_to_bdf(self.temp_dir,
                              bdf_path=str(Path(self.temp_dir, 'mydata.bdf').resolve()))

        self.assertEqual(Path(path).name, 'mydata.bdf')

    def test_convert_to_bdf_raises_value_error_with_invalid_edfpath(self):
        """Test that a value error is raised with the bdf path is invalid"""
        bdf_path = 'invalid.pdf'
        with self.assertRaises(ValueError):
            convert_to_bdf(self.temp_dir, bdf_path=bdf_path)

    def test_convert_to_bdf_with_write_targetness(self):
        """Test creating the EDF using targetness for event annotations"""
        path = convert_to_bdf(self.temp_dir,
                              write_targetness=True)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            edf = read_raw_bdf(path, preload=True)

        self.assertIn('target', edf.annotations.description)
        self.assertIn('nontarget', edf.annotations.description)

    def test_convert_to_bdf_without_write_targetness(self):
        """Test creating the EDF with labels as event annotations"""
        path = convert_to_bdf(self.temp_dir,
                              write_targetness=False)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            edf = read_raw_bdf(path, preload=True)

        self.assertNotIn('target', edf.annotations.description)
        self.assertNotIn('nontarget', edf.annotations.description)
        self.assertIn('+', edf.annotations.description)

    def test_convert_to_bdf_with_pre_filter(self):
        """Test creating the BDF with pre-filtering"""

        data = mock()
        channels = mock()
        fs = 100
        events = mock()
        annotation_channels = 1
        return_path = mock()
        when(convert).pyedf_convert(
            any_value(str),
            write_targetness=any_value(bool),
            use_event_durations=any_value(bool),
            remove_pre_fixation=any_value(bool),
            pre_filter=True,
        ).thenReturn((
            data,
            channels,
            fs,
            events,
            annotation_channels,
            'filter_settings'
        ))
        when(convert).write_pyedf(
            any_value(str),
            data,
            channels,
            fs,
            events,
            any_value(int),
            any_value(bool),
            annotation_channels,
            any_value(str)).thenReturn(return_path)
        path = convert_to_bdf(self.temp_dir,
                              pre_filter=True)

        verify(convert, times=1).pyedf_convert(
            any_value(str),
            write_targetness=any_value(bool),
            use_event_durations=any_value(bool),
            remove_pre_fixation=any_value(bool),
            pre_filter=True,
        )
        verify(convert, times=1).write_pyedf(
            any_value(str),
            data,
            channels,
            fs,
            events,
            any_value(int),
            any_value(bool),
            annotation_channels,
            any_value(str))
        self.assertEqual(path, return_path)


class TestPyedfconvert(unittest.TestCase):

    def setUp(self) -> None:
        self.temp_dir = tempfile.mkdtemp()
        self.sample_rate = 300
        self.filter_settings = {
            'filter_low': 0.5,
            'filter_high': 30,
            'filter_order': 5,
            'notch_filter_frequency': 60,
            'down_sampling_rate': 3,
            'static_trigger_offset': 0.0
        }
        create_bcipy_session_artifacts(
            self.temp_dir,
            sample_rate=self.sample_rate,
            filter_settings=self.filter_settings,
            samples=10000)  # sample number determines the annotation count, set it high here!

    def tearDown(self) -> None:
        shutil.rmtree(self.temp_dir)
        verifyNoMoreInteractions()
        unstub()

    def test_pyedf_convert_defaults(self):
        """Test the pyedf_convert function"""
        data, channels, fs, events, annotation_channels, pre_filter = pyedf_convert(
            self.temp_dir)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        self.assertFalse(pre_filter)

    def test_pyedf_convert_with_pre_filter(self):
        """Test the pyedf_convert function with pre-filtering"""
        data, channels, fs, events, annotation_channels, pre_filter = pyedf_convert(
            self.temp_dir, pre_filter=True)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate / self.filter_settings['down_sampling_rate'])
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        #  with less samples, the number of annotation channels increases
        self.assertEqual(annotation_channels, 2)
        self.assertIsInstance(pre_filter, str)

    def test_pyedf_convert_with_write_targetness(self):
        """Test the pyedf_convert function with targetness for event annotations"""
        data, channels, fs, events, annotation_channels, _ = pyedf_convert(
            self.temp_dir, write_targetness=True)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        for event in events:
            _, _, label = event
            # in this casae label and targetness should not be the same
            self.assertTrue(label in ['target', 'nontarget'])

    def test_pyedf_convert_without_write_targetness(self):
        """Test the pyedf_convert function with labels as event annotations"""
        data, channels, fs, events, annotation_channels, _ = pyedf_convert(
            self.temp_dir, write_targetness=False)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        for event in events:
            _, _, label = event
            self.assertTrue(label not in ['target', 'nontarget'])

    def test_pyedf_convert_with_use_event_durations(self):
        """Test the pyedf_convert function with event durations"""
        data, channels, fs, events, annotation_channels, _ = pyedf_convert(
            self.temp_dir, use_event_durations=True)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        for _, duration, _ in events:
            self.assertTrue(duration > 0)

    def test_pyedf_convert_without_use_event_durations(self):
        """Test the pyedf_convert function without event durations"""
        data, channels, fs, events, annotation_channels, _ = pyedf_convert(
            self.temp_dir, use_event_durations=False)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        for _, duration, _ in events:
            self.assertEqual(duration, 0)

    def test_pyedf_convert_with_remove_pre_fixation(self):
        """Test the pyedf_convert function with pre-fixation removal"""
        data, channels, fs, events, annotation_channels, _ = pyedf_convert(
            self.temp_dir, remove_pre_fixation=True)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(channels) > 0)
        self.assertEqual(fs, self.sample_rate)
        self.assertTrue(len(events) > 2)
        # see the validate_annotations function for details on how this is calculated
        self.assertEqual(annotation_channels, 1)
        for event in events:
            self.assertNotIn('fixation', event)
            self.assertNotIn('prompt', event)


class TestMNEConvert(unittest.TestCase):
    """Test the convert_to_mne function, which converts bcipy RawData into an mne format"""

    def setUp(self):
        """Set up the test case with a temporary directory and sample data"""
        self.temp_dir = tempfile.mkdtemp()
        self.sample_rate = 300
        self.filter_settings = {
            'filter_low': 0.5,
            'filter_high': 30,
            'filter_order': 5,
            'notch_filter_frequency': 60,
            'down_sampling_rate': 3,
            'static_trigger_offset': 0.0
        }
        self.channels = ['timestamp', 'O1', 'O2', 'Pz']
        self.raw_data = RawData('SampleDevice', self.sample_rate, self.channels)
        # generate 100 random samples of data
        for _ in range(0, 100):
            channel_data = gen_random_data(low=-1000,
                                           high=1000,
                                           channel_count=len(self.channels))
            self.raw_data.append(channel_data)

    def tearDown(self):
        """Remove the temporary directory and its contents after each test"""
        shutil.rmtree(self.temp_dir)

    def test_convert_to_mne_defaults(self):
        """Test the convert_to_mne function with default parameters"""
        data = convert_to_mne(self.raw_data)

        self.assertTrue(len(data) > 0)
        self.assertEqual(data.ch_names, self.channels[1:])
        self.assertEqual(data.info['sfreq'], self.sample_rate)

    def test_convert_to_mne_with_channel_map(self):
        """Test the convert_to_mne function with channel mapping"""
        # here we know only three channels are generated, using the channel map let's only use the last one
        channel_map = [0, 0, 1]
        data = convert_to_mne(self.raw_data, channel_map=channel_map)

        self.assertTrue(len(data) > 0)
        self.assertTrue(len(data.ch_names) == 1)  # this is the main assertion!
        self.assertEqual(data.info['sfreq'], self.sample_rate)

    def test_convert_to_mne_with_channel_types(self):
        """Test the convert_to_mne function with channel types"""
        channel_types = ['eeg', 'eeg', 'seeg']
        data = convert_to_mne(self.raw_data, channel_types=channel_types)

        self.assertTrue(len(data) > 0)
        self.assertEqual(data.ch_names, self.channels[1:])
        self.assertEqual(data.info['sfreq'], self.sample_rate)
        self.assertTrue(data.get_channel_types()[2] == 'seeg')

    def test_convert_to_mne_with_transform(self):
        """Test the convert_to_mne function with a transform"""
        multiplier = 2

        def transform(x, fs):
            return x * multiplier, fs

        data = convert_to_mne(self.raw_data, transform=transform)

        self.assertTrue(len(data) > 0)
        self.assertEqual(data.ch_names, self.channels[1:])
        self.assertEqual(data.info['sfreq'], self.sample_rate)

        # apply the transform to the first data point and compare to data returned
        expected_first_data_point = self.raw_data.channel_data[0][0] * multiplier
        self.assertTrue(data.get_data()[0][0] == expected_first_data_point)

    def test_convert_to_mne_with_mv_conversion(self):
        """Test the convert_to_mne function with a mv conversion"""
        data = convert_to_mne(self.raw_data, volts=False)

        self.assertTrue(len(data) > 0)
        self.assertEqual(data.ch_names, self.channels[1:])
        self.assertEqual(data.info['sfreq'], self.sample_rate)

        # apply the transform to the first data point and compare to data returned
        expected_first_data_point = self.raw_data.channel_data[0][0] * 1e-6
        self.assertTrue(data.get_data()[0][0] == expected_first_data_point)

    def test_convert_to_mne_with_custom_montage(self):
        """Test the convert_to_mne function with a custom montage"""

        # see https://mne.tools/stable/auto_tutorials/intro/40_sensor_locations.html
        # for more information on montages and available defaults
        montage_type = 'biosemi64'
        data = convert_to_mne(self.raw_data, montage=montage_type)

        self.assertTrue(len(data) > 0)
        self.assertEqual(data.ch_names, self.channels[1:])
        self.assertEqual(data.info['sfreq'], self.sample_rate)


class TestCompressionSupport(unittest.TestCase):

    def setUp(self):
        self.dir_name = 'test/'
        self.tar_file_name = 'test_file'
        self.tar_file_full_name = f'{self.tar_file_name}.tar.gz'
        self.test_file_name = 'test.text'
        with open(self.test_file_name, 'w', encoding=DEFAULT_ENCODING) as fp:
            pass

    def tearDown(self):
        os.remove(self.test_file_name)

        if os.path.exists(self.tar_file_full_name):
            os.remove(self.tar_file_full_name)

    def test_compression_writes_tar_gz_no_extension(self):
        # Test with no extension on tar name
        compress(self.tar_file_name, [self.test_file_name])
        # Assert correct file was written
        self.assertTrue(os.path.exists(self.tar_file_full_name))

    def test_compression_writes_tar_gz_with_extension(self):
        # Test with extension on tar name
        compress(self.tar_file_full_name, [self.test_file_name])
        # Assert correct file was written
        self.assertTrue(os.path.exists(self.tar_file_full_name))

    def test_decompression_extracts_file_no_extension(self):
        # Test with no extension on tar name
        compress(self.tar_file_name, [self.test_file_name])
        decompress(self.tar_file_name, ".")
        self.assertTrue(os.path.exists(self.test_file_name))

    def test_decompression_extracts_file_with_extension(self):
        # Test with extension on tar name
        compress(self.tar_file_name, [self.test_file_name])
        decompress(self.tar_file_full_name, ".")
        self.assertTrue(os.path.exists(self.test_file_name))

    def test_file_not_found_error_thrown_on_compression(self):
        garbage_name = 'not_possible_to_exist.biz'
        with self.assertRaises(FileNotFoundError):
            compress(self.tar_file_name, [garbage_name])

    def test_file_list_returns_compressed_file_name_no_extension(self):
        # Test with no extension on tar name
        compress(self.tar_file_name, [self.test_file_name])
        tar_list = archive_list(self.tar_file_name)
        self.assertTrue(tar_list[0] == self.test_file_name)

    def test_file_list_returns_compressed_file_name_with_extension(self):
        # Test with extension on tar name
        compress(self.tar_file_name, [self.test_file_name])
        tar_list = archive_list(self.tar_file_full_name)
        self.assertTrue(tar_list[0] == self.test_file_name)


class TestConvertTobii(unittest.TestCase):

    def test_tobii_to_norm(self):
        """Test the tobii_to_norm function"""
        tobii_data = (0.5, 0.5)  # center of screen in tobii coordinates
        excepted_norm_data = (0, 0)  # center of screen in norm coordinates
        norm_data = tobii_to_norm(tobii_data)
        self.assertEqual(norm_data, excepted_norm_data)

        tobii_data = (0, 0)  # top left of screen in tobii coordinates
        excepted_norm_data = (-1, 1)  # top left of screen in norm coordinates
        norm_data = tobii_to_norm(tobii_data)
        self.assertEqual(norm_data, excepted_norm_data)

        tobii_data = (1, 1)  # bottom right of screen in tobii coordinates
        excepted_norm_data = (1, -1)  # bottom right of screen in norm coordinates
        norm_data = tobii_to_norm(tobii_data)
        self.assertEqual(norm_data, excepted_norm_data)

    def test_tobii_to_norm_raises_error_with_invalid_units(self):
        """Test the tobii_to_norm function raises an error with invalid units"""
        tobii_data = (-1, 1)  # invalid tobii coordinates
        with self.assertRaises(AssertionError):
            tobii_to_norm(tobii_data)

        tobii_data = (1, 11)  # invalid tobii coordinates

        with self.assertRaises(AssertionError):
            tobii_to_norm(tobii_data)

    def test_norm_to_tobii(self):
        """Test the norm_to_tobii function"""
        norm_data = (0, 0)  # center of screen in norm coordinates
        excepted_tobii_data = (0.5, 0.5)  # center of screen in tobii coordinates
        tobii_data = norm_to_tobii(norm_data)
        self.assertEqual(tobii_data, excepted_tobii_data)

        norm_data = (-1, 1)  # top left of screen in norm coordinates
        excepted_tobii_data = (0, 0)  # top left of screen in tobii coordinates
        tobii_data = norm_to_tobii(norm_data)
        self.assertEqual(tobii_data, excepted_tobii_data)

        norm_data = (1, -1)  # bottom right of screen in norm coordinates
        excepted_tobii_data = (1, 1)  # bottom right of screen in tobii coordinates
        tobii_data = norm_to_tobii(norm_data)
        self.assertEqual(tobii_data, excepted_tobii_data)

    def test_norm_to_tobii_raises_error_with_invalid_units(self):
        """Test the norm_to_tobii function raises an error with invalid units"""
        norm_data = (-1.1, 1)
        with self.assertRaises(AssertionError):
            norm_to_tobii(norm_data)

        norm_data = (1, 1.1)
        with self.assertRaises(AssertionError):
            norm_to_tobii(norm_data)


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