from ada.io.file_manager import FileManager
from ada.preprocessing.epocher import ActivityIndex
from ada.io.acti_eeg import ActiPSG
from ada.data_containers.scored import ScoredShort
from ada.short.cole_kripke import ColeKripke, Sazonov, Ucsd, Webster, Scripps
from ada.short.unified import UFA
from ada.data_containers._tag_reader import TagsFileReader
import numpy as np
import os
from unittest import TestCase
import pytest
from math import isclose


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


@pytest.mark.order(6)
class TestScoringShort:

    raw = FileManager.load_file(os.path.join(DIRECTORY, 'raw.ada'))
    epoched_30 = ActivityIndex(30).to_epoch(raw)
    epoched_60 = ActivityIndex(60).to_epoch(raw)
    thresholds = {ColeKripke: 2e-4,
                  Sazonov: 0.4,
                  Ucsd: 1.5,
                  Webster: 0.4,
                  Scripps: 1.2}
    tags_30 = TagsFileReader(os.path.join(DIRECTORY, 'raw.obci.tag')).get_tags()
    tags_60 = tags_30[::2]

    def get_test_signals(self, scorer):
        if scorer._input_epoch == 30:
            default = self.epoched_30
            wrong = self.epoched_60
        elif scorer._input_epoch == 60:
            default = self.epoched_60
            wrong = self.epoched_30

        return default, wrong
    
    def scoring_helper(self, scorer):
        default, wrong = self.get_test_signals(scorer(1))

        # with rescoring if present in algorithm, or when not present
        if isinstance(scorer.__init__.__defaults__[-1], bool):
            assert scorer.__init__.__defaults__[-1]
        
        score = scorer(self.thresholds[scorer]).to_score(default)
        loaded = ScoredShort.load_file(os.path.join(DIRECTORY, scorer.__name__ + '.ada'))
        assert isinstance(score, ScoredShort)
        assert np.all(score.data == loaded.data)
        assert score.first_sample_timestamp == loaded.first_sample_timestamp
        assert score.last_sample_timestamp == loaded.last_sample_timestamp
        assert score.channel_names == loaded.channel_names
        assert score.fs == loaded.fs
        TestCase().assertDictEqual(score.scoring_method_metadata, loaded.scoring_method_metadata)
        TestCase().assertDictEqual(score.epoching_method_metadata, loaded.epoching_method_metadata)
        TestCase().assertDictEqual(score.metadata, loaded.metadata)
        with pytest.raises(ValueError):
            _ = scorer(1).to_score(wrong)

        # without rescoring if present
        if scorer.__init__.__code__.co_argcount == 3:
            score = scorer(self.thresholds[scorer], False).to_score(default)
            loaded = ScoredShort.load_file(os.path.join(DIRECTORY, scorer.__name__ + '_no_rescoring.ada'))
            assert np.all(score.data == loaded.data)

    def threshold_helper(self, scorer, corr, threshold):
        default, _ = self.get_test_signals(scorer(1))
        if default.fs == 1 / 60:
            tags = self.tags_60
        else:
            tags = self.tags_30
        
        acti_psg_1 = ActiPSG.from_continous_stages(default, tags, 1711512330.6438127, trim_to_eeg=True)
        s = scorer(1)
        mcc = s.fit_threshold([acti_psg_1], (self.thresholds[scorer] / 10, self.thresholds[scorer] * 10))
        assert isclose(threshold, s._threshold)
        assert isclose(mcc, corr)

    def interpolation_helper(self, scorer, threshold):
        default, _ = self.get_test_signals(scorer(1))
        s = scorer(threshold=None)
        t = s._interpolate_threshold(default)
        assert isclose(t, threshold)

    
    def test_ColeKripke(self):
        self.scoring_helper(ColeKripke)
        self.threshold_helper(ColeKripke, 0.7947194142390263, 3.1891891891891894e-05)
        self.interpolation_helper(ColeKripke, 3.5225225225225225e-5)

    
    def test_Sazonov(self):
        self.scoring_helper(Sazonov)
        self.threshold_helper(Sazonov, 0.05971453032500155, 0.20252252252252254)
        self.interpolation_helper(Sazonov, 0.26361842122537515)


    def test_Ucsd(self):
        self.scoring_helper(Ucsd)
        self.threshold_helper(Ucsd, 0.6111111111111112, 0.7)
        self.interpolation_helper(Ucsd, 0.8308308308308309)


    def test_Webster(self):
        self.scoring_helper(Webster)
        self.threshold_helper(Webster, 0.5477225575051662, 0.15891891891891893)
        self.interpolation_helper(Webster, 0.14904904904904906)


    def test_Scripps(self):
        self.scoring_helper(Scripps)
        self.threshold_helper(Scripps, 0.16461664267663798, 11.227027027027026)
        self.interpolation_helper(Scripps, 1.1789815490718634)

    def test_UFA(self):
        score = UFA(0.75).to_score(self.epoched_30)
        loaded = ScoredShort.load_file(os.path.join(DIRECTORY, 'UFA.ada'))
        assert isinstance(score, ScoredShort)
        assert np.all(score.data == loaded.data)
        assert score.first_sample_timestamp == loaded.first_sample_timestamp
        assert score.last_sample_timestamp == loaded.last_sample_timestamp
        assert score.channel_names == loaded.channel_names
        assert score.fs == loaded.fs
        TestCase().assertDictEqual(score.scoring_method_metadata, loaded.scoring_method_metadata)
        TestCase().assertDictEqual(score.epoching_method_metadata, loaded.epoching_method_metadata)
        TestCase().assertDictEqual(score.metadata, loaded.metadata)
        
        acti_psg_1 = ActiPSG.from_continous_stages(self.epoched_30, self.tags_30, 1711512330.6438127, trim_to_eeg=True)
        s = UFA(1)
        mcc = s.fit_threshold([acti_psg_1], (0.075, 7.5))
        assert isclose(1.3459459459459457, s._threshold)
        assert isclose(mcc, 0.46080494552700013)

        threshold = s._interpolate_threshold(self.epoched_30)
        assert isclose(threshold, 0.3424948234307415)


@pytest.mark.order(7)
class TestShortMetrics():
    temp = np.empty((2, 20))
    temp[0] = [1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 1, 1]
    temp[1] = np.arange(0, 60 * 20, 60)
    data = ScoredShort(temp, {}, 1 / 60, {}, {})

    def test_total_sleep_time(self):
        assert self.data.total_sleep_time() == 12.
    
    def test_sleep_efficiency(self):
        assert self.data.sleep_efficiency() == 12 / 20 * 100

    def test_sleep_onset(self):
        assert self.data.sleep_onset(5) == 2.

    def test_sleep_episodes(self):
        assert self.data.sleep_episodes() == 4

    def test_mean_sleep_episode(self):
        assert self.data.mean_sleep_episode() == 3.

    def test_long_sleep_episodes(self):
        assert self.data

    def test_sleep_fragmentation_index(self):
        assert self.data.sleep_fragmentation_index() == 5 / 12

    def test_waso(self):
        assert self.data.wake_after_sleep_onset(5) == 4.
