from ada.io.file_manager import FileManager
from ada.preprocessing.epocher import ActivityIndex, MVM, Downsampler, CKEpocher, MIMS
from ada.io.acti_eeg import ActiPSG
import numpy as np
import os
from unittest import TestCase
import pytest


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


@pytest.mark.order(4)
class TestEpochingActi:

    raw = FileManager.load_file(os.path.join(DIRECTORY_IN, 'epochingExample.ada'))

    def epoching_helper(self, epocher, raw):
        epocher = epocher(1)
        epoched = epocher.to_epoch(raw)
        epoched_test = FileManager.load_file(os.path.join(DIRECTORY, type(epocher).__name__ + '.ada'))
        assert epoched.fs == 1.
        assert epoched.channel_names == epoched_test.channel_names
        assert np.allclose(epoched.data, epoched_test.data)
        TestCase().assertDictEqual(epoched.metadata, epoched_test.metadata)
        TestCase().assertDictEqual(epoched_test.epoching_method_metadata, epoched.epoching_method_metadata)


    def test_ActivityIndex(self):
        self.epoching_helper(ActivityIndex, self.raw)

    
    def test_CKEpocher(self):
        self.epoching_helper(CKEpocher, self.raw)


    def test_MVM(self):
        self.epoching_helper(MVM, self.raw)


    def test_Downsampler(self):
        self.epoching_helper(Downsampler, self.raw)

    
    def test_MIMS(self):  # TODO Stage 2 of MIMS is not tested (no maxed out data)
        self.epoching_helper(MIMS, self.raw)


@pytest.mark.order(5)
class TestEpochingActiPsg:
    def test_epoching(self):
        raw = FileManager.load_file(os.path.join(DIRECTORY_IN, 'epochingExample.ada'))
        tags = [dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['s1', 1706472231.5 + 100 + i * 20, 1706472231.5 + 100 + (i + 1) * 20])) for i in range(1200 // 20)]
        tags[20]['name'] = 'w'
        acti_psg = ActiPSG.from_continous_stages(raw, tags, 1706472231.5)
        
        # long epoch
        epocher = ActivityIndex(60)
        epoched = epocher.to_epoch(acti_psg, 0.3)
        epoched_acti = epocher.to_epoch(acti_psg.acti_data)
        assert epoched.acti_data.fs == 1 / 60
        assert isinstance(epoched, ActiPSG)
        assert np.allclose(epoched.acti_data.data, epoched_acti.data)
        assert np.isclose(epoched.first_sample_timestamp, acti_psg.first_sample_timestamp)
        new_stages = np.sum(np.array(np.split(acti_psg.eeg_data.psg_stages, 25)), axis=1) >= 75 * 20
        assert np.all(epoched.eeg_data.psg_stages == new_stages)

        # short epoch, not with 1/length equal to integer divisor of original fs
        epocher = ActivityIndex(1 / 15)
        epoched = epocher.to_epoch(acti_psg, 0.3)
        epoched_acti = epocher.to_epoch(acti_psg.acti_data)
        assert epoched.acti_data.fs == 15
        assert isinstance(epoched, ActiPSG)
        assert np.allclose(epoched.acti_data.data, epoched_acti.data)
        assert np.isclose(epoched.first_sample_timestamp, acti_psg.first_sample_timestamp)
        new_stages = np.sum(np.array(np.split(acti_psg.eeg_data.psg_stages, 22500)), axis=1) >= 1.5
        assert np.all(epoched.eeg_data.psg_stages == new_stages)

        # TODO right now it does not work, if 1/length is not integer divisor of fs
