from ada.io.file_manager import FileManager
from ada.data_containers.generic import GenericData
from unittest import TestCase
import numpy as np
import os
import pytest


DIRECTORY = os.path.join(os.path.dirname(__file__), 'test_reading_data_files')


@pytest.mark.order(2)
class TestCuttingActi:
    def test_cutting_raw_geneactiv(self):
        raw = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBin.bin'))
        cutted = raw.cut_by_timestamp(raw.first_sample_timestamp + 0.5, raw.first_sample_timestamp + 2.5)
        assert np.all(raw.data[:, 50:250] == cutted.data)
        cutted = raw.cut_by_samples(50, 250)
        assert np.all(raw.data[:, 50:250] == cutted.data)
        assert raw.fs == cutted.fs
        assert raw.channel_names == cutted.channel_names
        assert isinstance(cutted, type(raw))
        TestCase().assertDictEqual(raw.metadata, cutted.metadata)


    def test_cutting_raw_generic(self):
        raw = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBin.bin'))
        generic_raw = GenericData.from_nongeneric(raw)
        generic_cutted = generic_raw.cut_by_samples(50, 250)
        assert np.all(generic_raw.data[:, 50:250] == generic_cutted.data)
        generic_cutted = generic_raw.cut_by_timestamp(raw.first_sample_timestamp + 0.5, raw.first_sample_timestamp + 2.5)
        assert np.all(generic_raw.data[:, 50:250] == generic_cutted.data)
        assert generic_raw.fs == generic_cutted.fs
        assert generic_raw.channel_names == generic_cutted.channel_names
        assert isinstance(generic_cutted, type(generic_raw))
        TestCase().assertDictEqual(generic_raw.metadata, generic_cutted.metadata)
        assert generic_raw.epoching_method_metadata == generic_cutted.epoching_method_metadata


    def test_cutting_epoched_geneactiv(self):
        epoched = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        cutted = epoched.cut_by_samples(1, 6)
        assert np.all(cutted.data == epoched.data[:, 1:6])
        cutted = epoched.cut_by_timestamp(epoched.first_sample_timestamp + 1, epoched.first_sample_timestamp + 301)
        assert np.all(cutted.data == epoched.data[:, 1:6])
        assert epoched.fs == cutted.fs
        assert epoched.channel_names == cutted.channel_names
        assert isinstance(cutted, type(epoched))
        TestCase().assertDictEqual(epoched.metadata, cutted.metadata)
        TestCase().assertDictEqual(epoched.epoching_method_metadata, cutted.epoching_method_metadata)


    def test_cutting_epoched_generic(self):
        epoched = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        generic_epoched = GenericData.from_nongeneric(epoched)
        generic_cut = generic_epoched.cut_by_samples(1, 6)
        assert np.all(generic_cut.data == generic_epoched.data[:, 1:6])
        generic_cut = generic_epoched.cut_by_timestamp(generic_epoched.first_sample_timestamp + 1, generic_epoched.first_sample_timestamp + 301)
        assert np.all(generic_cut.data == generic_epoched.data[:, 1:6])
        assert generic_epoched.fs == generic_cut.fs
        assert generic_epoched.channel_names == generic_cut.channel_names
        assert isinstance(generic_cut, type(generic_epoched))
        TestCase().assertDictEqual(generic_epoched.metadata, generic_cut.metadata)
        TestCase().assertDictEqual(generic_epoched.epoching_method_metadata, generic_cut.epoching_method_metadata)

    
    def test_cutting_mesa(self):
        mesa = FileManager.load_file(os.path.join(DIRECTORY, 'MESA.csv'))
        mesa_cut = mesa.cut_by_samples(1, 3)
        assert np.allclose(mesa.data[:, 1:3], mesa_cut.data, equal_nan=True)
        assert mesa.fs == mesa_cut.fs
        assert mesa.channel_names == mesa_cut.channel_names
        assert isinstance(mesa_cut, type(mesa))
        metadata = mesa.metadata.copy()
        metadata['first_line'] = mesa.metadata['first_line'] + 1
        TestCase().assertDictEqual(metadata, mesa_cut.metadata)
        TestCase().assertDictEqual(mesa.epoching_method_metadata, mesa_cut.epoching_method_metadata)
        TestCase().assertDictEqual(mesa.scoring_method_metadata, mesa_cut.scoring_method_metadata)
        mesa_cut = mesa.cut_by_timestamp(1, 70)
        assert np.allclose(mesa.data[:, 1:3], mesa_cut.data, equal_nan=True)


    def test_cutting_actigraph(self):
        data = FileManager.load_file(os.path.join(DIRECTORY, 'neo.gt3x'))
        data_cut = data.cut_by_samples(1, 10)
        assert np.allclose(data.data[1:10], data_cut.data)
        assert data.fs == data_cut.fs
        assert isinstance(data_cut, type(data))
        TestCase().assertDictEqual(data.metadata, data_cut.metadata)
        assert data_cut.channel_names == data.channel_names
        data_cut_ts = data.cut_by_timestamp(637755981.0333333, 637755981.3333334)
        assert np.allclose(data_cut_ts.data, data_cut.data)


@pytest.mark.order(3)
class TestCuttingActiPsg:
    def test_cutting(self):
        acti_psg = FileManager.load_file(os.path.join(DIRECTORY, 'ActiPSG.zip'))
        start_timestamp = acti_psg.acti_data.first_sample_timestamp
        end_timestamp = acti_psg.acti_data.last_sample_timestamp
        cut = acti_psg.cut_by_timestamp(start_timestamp + 59.123, end_timestamp - 14.876)
        assert isinstance(cut, type(acti_psg))
        assert cut.acti_data.fs == acti_psg.acti_data.fs
        assert cut.acti_data.channel_names == acti_psg.acti_data.channel_names
        assert np.isclose(cut.acti_data.first_sample_timestamp, start_timestamp + 59.123)
        assert np.isclose(cut.acti_data.last_sample_timestamp, end_timestamp - 14.876)
        assert np.isclose(cut.eeg_data.start_timestamp, cut.acti_data.first_sample_timestamp)
        assert np.isclose(cut.eeg_data.end_timestamp, cut.acti_data.last_sample_timestamp)
        assert np.all(cut.acti_data.data == acti_psg.acti_data.data[:, int(59.123 / 30) + 1:-1])
