from ada.io.file_manager import FileManager
from ada.long.cosinor import Linear, Nonlinear, Antilogistic, Arctangent, Saw, Hill, CosinorResults
from ada.long.ar import Spectrum, ARResults
from ada.long.nonparametric import DFA, DFAResults, SimpleNonparametric, SimpleNonparametricResults

import numpy as np
import os
import pytest
from math import isclose
import scipy.signal as ss


DIRECTORY = os.path.join(os.path.dirname(__file__), 'test_scoring_data_files')
DATA = FileManager.load_file(os.path.join(DIRECTORY, 'long.ada'))


def _dict_isclose(dict1, dict2, rtol=1e-9):
    assert dict1.keys() == dict2.keys()
    for key in dict1.keys():
        if isinstance(dict1[key], dict):
            _dict_isclose(dict1[key], dict2[key], rtol)
        elif isinstance(dict1[key], (float, complex)):
            assert isclose(dict1[key], dict2[key], rel_tol=rtol)
        elif isinstance(dict1[key], (tuple, list, np.ndarray)):
            assert len(dict1[key]) == len(dict2[key])
            for e1, e2 in zip(dict1[key], dict2[key]):
                if isinstance(e1, (float, complex)):
                    assert isclose(e1, e2, rel_tol=rtol)
                else:
                    assert e1 == e2
        else:
            assert dict1[key] == dict2[key]


@pytest.mark.order(8)
class TestCosinor:

    def scoring_helper(self, algorithm, results, *args):
        scorer = algorithm(*args)
        out = scorer.fit(DATA, None)
        for i, r in enumerate(results):
            assert out[i].id == 'P78'
            _dict_isclose(out[i].fit_params, r, 1e-7)
    

    def test_data_getter(self):  # this is defined in abstract overclass, so we can use any class as a test for all
        data = FileManager.load_file(os.path.join(DIRECTORY, 'raw.ada'))
        t, to_score = Linear._get_data(data, None)
        assert np.allclose(data.timestamp, t)
        assert np.allclose(data.to_score, to_score)
        t, x = Linear._get_data(data, 'x')
        assert np.allclose(data.timestamp, t)
        assert np.allclose(data.x, x)


    def test_preprocessing(self):  # this is defined in abstract overclass, so we can use any class as a test for all
        sos = ss.iirdesign(1 / (3600 * 9), 1 / (3600 * 3), 2, 30, ftype='butter', output='sos', fs=DATA.fs)
        out = ss.sosfiltfilt(sos, DATA.to_score - np.mean(DATA.to_score))
        preprocessed = Linear._preprocess(DATA.to_score, DATA.fs)
        assert np.allclose(out, preprocessed)

    
    def test_linear(self):
        self.scoring_helper(Linear, [{'base freq': {'period': 24.0, 'amplitude': 4.248651770853655, 'phase': 0.935340078488248},
                                      'amplitude': 4.2486436720586305,
                                      'constant': 0.0485283214148482,
                                      'f-test': 0.0,
                                      'bootstrap': 1.1102230246251565e-16,
                                      'method': 'Linear',
                                      'preprocessing': True,
                                      'channel': 'default',
                                      'id': 'P78'}],
                            24, True, True, 1)
        
        self.scoring_helper(Linear, [{'base freq': {'period': 24.0, 'amplitude': 2.5870534777740373, 'phase': -0.2333232934186098},
                                      '1 harmonic': {'period': 12.0, 'amplitude': 3.471677884860785, 'phase': 1.2244199099077768},
                                      '2 harmonic': {'period': 8.0, 'amplitude': 0.8706753091558049, 'phase': -0.292070763147783},
                                      'amplitude': 4.782262109319531,
                                      'constant': 6.549395413531882,
                                      'f-test': 0.0,
                                      'bootstrap': None,
                                      'method':
                                      'Linear',
                                      'preprocessing': False,
                                      'channel': 'default',
                                      'id': 'P78'},
                                     {'base freq': {'period': 26.0, 'amplitude': 2.229147646329575, 'phase': 3.2904745421315678},
                                      '1 harmonic': {'period': 13.0, 'amplitude': 1.4333454297313222, 'phase': 1.0640488012316567},
                                      '2 harmonic': {'period': 8.666666666666666, 'amplitude': 0.9317289083160062, 'phase': 2.2332575254699707},
                                      'amplitude': 3.2535346212272684,
                                      'constant': 6.671901105174544,
                                      'f-test': 4.583556637684447e-154,
                                      'bootstrap': None,
                                      'method': 'Linear',
                                      'preprocessing': False,
                                      'channel': 'default',
                                      'id': 'P78'}],
                                    [24, 26], False, False, 3)
        
        # check case which shouldn't be significant
        self.scoring_helper(Linear, [{'base freq': {'period': 40.0, 'amplitude': 0.049523054146070035, 'phase': -0.392218141174852},
                                      'amplitude': 0.04952304204387037,
                                      'constant': 6.574717242982166,
                                      'f-test': 0.899363639958302,
                                      'bootstrap': 0.19478902251089059,
                                      'method': 'Linear', 'preprocessing': False,
                                      'channel': 'default',
                                      'id': 'P78'}],
                            40, False, True, 1)
    
        
    def test_nonlinar(self):
        self.scoring_helper(Nonlinear, [{'base freq': {'period': 24.050052536504495, 'amplitude': 4.129204955075128, 'phase': 0.9618226244880024},
                                         '1 harmonic': {'period': 12.025026268252248, 'amplitude': 1.0307654185255182, 'phase': 2.2924617218845493},
                                         '2 harmonic': {'period': 8.016684178834831, 'amplitude': 1.2062898626491783, 'phase': -0.12268344516509178},
                                         'amplitude': 5.460743014875334,
                                         'constant': 6.558686248519991,
                                         'f-test': 0.0,
                                         'bootstrap': None,
                                         'method': 'Nonlinear',
                                         'preprocessing': False,
                                         'channel': 'default',
                                         'id': 'P78'}],
                            (16, 32), False, False, 3)
        
        self.scoring_helper(Nonlinear, [{'base freq': {'period': 24.087042268097676, 'amplitude': 4.257492889020228, 'phase': 1.016052832718863},
                                         'amplitude': 4.2574928604663,
                                         'constant': 0.06184317516793147,
                                         'f-test': 0.0,
                                         'bootstrap': None,
                                         'method': 'Nonlinear',
                                         'preprocessing': True,
                                         'channel': 'default',
                                         'id': 'P78'}],
                            (16, 32), True, False, 1)


    def test_saw(self):
        self.scoring_helper(Saw, [{'amplitude': 5.388870329425846,
                                   'base freq': {'amplitude': 5.388870329425846, 'period': 24.085248159159118, 'phase': 2.2178030023424724, 'ratio': 0.37733051495605635},
                                   'bootstrap': None,
                                   'channel': 'default',
                                   'constant': 0.058701350613309715,
                                   'f-test': 0.0,
                                   'method': 'Saw',
                                   'preprocessing': True,
                                   'id': 'P78'}],
                            (16, 32), True, False)
        
    
    def test_antilogistic(self):
        self.scoring_helper(Antilogistic, [{'amplitude': -27.073630245620567,
                                            'base freq': {'alpha': 0.9999999999999999, 'amplitude': -27.073630245620567, 'beta': 0.7317305218582012, 'period': 24.08100341522575, 'phase': 2.549996838356318},
                                            'bootstrap': None,
                                            'channel': 'default',
                                            'constant': 9.119282354426362,
                                            'f-test': 0.0,
                                            'method': 'Antilogistic',
                                            'preprocessing': True,
                                            'id': 'P78'}],
                            (16, 32), True, False)
        

    def test_arctangent(self):
        self.scoring_helper(Arctangent, [{'amplitude': -36.29169271360979,
                                          'base freq': {'alpha': 0.9999999999999999, 'amplitude': -36.29169271360979, 'beta': 0.4491406285625499, 'period': 24.080821603459505, 'phase': 2.542662193139204},
                                          'bootstrap': None,
                                          'channel': 'default',
                                          'constant': -4.9813042501353175,
                                          'f-test': 0.0,
                                          'method': 'Arctangent',
                                          'preprocessing': True,
                                          'id': 'P78'}],
                            (16, 32), True, False)


    def test_hill(self):
        self.scoring_helper(Hill, [{'amplitude': -10.289460905326779,
                                    'base freq': {'amplitude': -10.289460905326779, 'gamma': 1.6435664809520893, 'm': 0.9999999999999999, 'period': 24.065700128418197, 'phase': 2.6398304322540445},
                                    'bootstrap': None,
                                    'channel': 'default',
                                    'constant': 4.435504356918882,
                                    'f-test': 0.0,
                                    'method': 'Hill',
                                    'preprocessing': True,
                                    'id': 'P78'}],
                            (16, 32), True, False)
        
    def test_loading(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'cosinor.ada.long'))
        assert isinstance(loaded, CosinorResults)
        _dict_isclose(loaded.fit_params, {'base freq': {'period': 24.0, 'amplitude': 4.248651770853655, 'phase': 0.935340078488248},
                                          'amplitude': 4.2486436720586305,
                                          'constant': 0.0485283214148482,
                                          'f-test': 0.0,
                                          'bootstrap': None,
                                          'method': 'Linear',
                                          'preprocessing': True,
                                          'channel': 'default',
                                          'id': 'P78'}, rtol=1e-7)
    
    def test_exporting(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'cosinor.ada.long'))
        loaded.export(os.path.join(DIRECTORY, 'temp_cosinor.ada.long'))
        exported = FileManager.load_file(os.path.join(DIRECTORY, 'temp_cosinor.ada.long'))
        assert isinstance(exported, CosinorResults)
        _dict_isclose(loaded.fit_params, exported.fit_params, rtol=1e-7)
        os.remove(os.path.join(DIRECTORY, 'temp_cosinor.ada.long'))


@pytest.mark.order(9)
class TestNonparametric:
    def test_dfa(self):
        scorer = DFA(n_crossovers=2)
        results = scorer.fit(DATA)
        _dict_isclose(results.fit_params, {'id': 'P78',
                                           'exponent': 0.9726391526417634,
                                           'exponent error': 0.004008683835462042,
                                           'scale range': (0.16666666666666666, 48.0),
                                           'trend degree': 1,
                                           'segments': {'segment 1': {'exponent': 0.9289578050683511, 'exponent_error': 0.004983129449502565},
                                                        'segment 2': {'exponent': 1.0810226549029465, 'exponent_error': 0.007972657155529283},
                                                        'segment 3': {'exponent': 0.08425523314071903, 'exponent_error': 0.054983026318238704},
                                                        'crossover 1': {'crossover': 2.944675898392218, 'ci': (2.4805933356379533, 3.495581489314128)},
                                                        'crossover 2': {'crossover': 26.84940260517049, 'ci': (25.740247217958277, 28.00635184853975)}}}, rtol=1e-7)
        a, b = results.crossovers
        assert np.allclose(a, [2.944675898392218, 26.84940260517049], rtol=1e-7)
        assert np.allclose(b, [(2.4805933356379533, 3.495581489314128), (25.740247217958277, 28.00635184853975)], rtol=1e-7)
        a, b = results.exponents
        assert np.allclose(a, [0.9289578050683511, 1.0810226549029465, 0.08425523314071903], rtol=1e-7)
        assert np.allclose(b, [0.004983129449502565, 0.007972657155529283, 0.054983026318238704], rtol=1e-7)
        assert results.id == 'P78'

    def test_dfa_unknown_crossovers(self):
        scorer = DFA(n_crossovers=None)
        results = scorer.fit(DATA)
        assert results.id == 'P78'
        _dict_isclose(results.fit_params, {'id': 'P78',
                                           'exponent': 0.9726391526417634,
                                           'exponent error': 0.004008683835462042,
                                           'scale range': (0.16666666666666666, 48.0),
                                           'trend degree': 1,
                                           'segments': {'segment 1': {'exponent': 0.8952935733817164, 'exponent_error': 0.010137951923644393},
                                                        'segment 2': {'exponent': 0.985736571672099, 'exponent_error': 0.008493231846618711},
                                                        'segment 3': {'exponent': 1.131781915647996, 'exponent_error': 0.013422286824701653},
                                                        'segment 4': {'exponent': 0.6938245102934644, 'exponent_error': 0.05856358290842304},
                                                        'segment 5': {'exponent': -1.2608388281949512, 'exponent_error': 0.40745798052886484},
                                                        'segment 6': {'exponent': 0.6041510251683007, 'exponent_error': 0.2434520585067723},
                                                        'crossover 1': {'crossover': 0.862104305872118, 'ci': (0.6439247682152182, 1.1542091108923447)},
                                                        'crossover 2': {'crossover': 6.021498375940666, 'ci': (5.012078903835101, 7.2342122674310145)},
                                                        'crossover 3': {'crossover': 20.648174318309582, 'ci': (18.832159995021133, 22.639309712322834)},
                                                        'crossover 4': {'crossover': 35.08366690089973, 'ci': (33.910905063260074, 36.29698708769751)},
                                                        'crossover 5': {'crossover': 39.77662283282746, 'ci': (38.001705123530066, 41.63444031897794)}}}, rtol=1e-7)

    def test_loading_dfa(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'dfa.ada.long'))
        assert isinstance(loaded, DFAResults)
        _dict_isclose(loaded.fit_params, {'id': 'P78',
                                          'exponent': 0.9726391526417634,
                                          'exponent error': 0.004008683835462042,
                                          'scale range': (0.16666666666666666, 48.0),
                                          'trend degree': 1,
                                          'segments': {'segment 1': {'exponent': 0.9289578050683511, 'exponent_error': 0.004983129449502565},
                                                       'segment 2': {'exponent': 1.0810226549029465, 'exponent_error': 0.007972657155529283},
                                                       'segment 3': {'exponent': 0.08425523314071903, 'exponent_error': 0.054983026318238704},
                                                       'crossover 1': {'crossover': 2.944675898392218, 'ci': (2.4805933356379533, 3.495581489314128)},
                                                       'crossover 2': {'crossover': 26.84940260517049, 'ci': (25.740247217958277, 28.00635184853975)}}}, rtol=1e-7)
        
    def test_exporting_dfa(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'dfa.ada.long'))
        loaded.export(os.path.join(DIRECTORY, 'temp_dfa.ada.long'))
        exported = FileManager.load_file(os.path.join(DIRECTORY, 'temp_dfa.ada.long'))
        assert isinstance(exported, DFAResults)
        _dict_isclose(loaded.fit_params, exported.fit_params, rtol=1e-7)
        os.remove(os.path.join(DIRECTORY, 'temp_dfa.ada.long'))


    def test_simple_nonparametric(self):
        scorer = SimpleNonparametric()
        results = scorer.fit(DATA)
        _dict_isclose(results.fit_params, {'IS': 0.488338385064575, 'IV': 0.7568962341885255, 'M10': 9.685750721212573, 'L5': 1.6802307622392285, 'id': 'P78'}, rtol=1e-7)
        assert results.id == 'P78'

    
    def test_loading_simple_nonparametric(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'snp.ada.long'))
        assert isinstance(loaded, SimpleNonparametricResults)
        _dict_isclose(loaded.fit_params, {'IS': 0.488338385064575, 'IV': 0.7568962341885255, 'M10': 9.685750721212573, 'L5': 1.6802307622392285, 'id': 'P78'}, rtol=1e-7)


    def test_exporting_simple_nonparametric(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'snp.ada.long'))
        loaded.export(os.path.join(DIRECTORY, 'temp_snp.ada.long'))
        exported = FileManager.load_file(os.path.join(DIRECTORY, 'temp_snp.ada.long'))
        assert isinstance(exported, SimpleNonparametricResults)
        _dict_isclose(loaded.fit_params, exported.fit_params, rtol=1e-7)
        os.remove(os.path.join(DIRECTORY, 'temp_snp.ada.long'))


@pytest.mark.order(10)
class TestAR:
    data = FileManager.load_file(os.path.join(DIRECTORY, 'longDownsampler.ada'))
    
    def test_spectrum(self):
        scorer = Spectrum(None, True, True)
        results = scorer.fit(self.data)
        temp_dict = dict([(k, results.fit_params[k]) for k in results.fit_params.keys() if k != 'coefficients'])
        _dict_isclose(temp_dict, {'noise std': 0.012509705347830354,
                                  'order': 180,
                                  'peak': 23.045592471672744,
                                  'preprocessing': True,
                                  'period ci': (23.04205230615615, 23.093883551649984), 'id': 'P66'}, rtol=1e-7)
        assert isclose(np.mean(results.freq), 2.0833333333333333e-05, rel_tol=1e7)
        assert isclose(np.mean(results.spectrum), 87.39790451480995, rel_tol=1e7)
        assert isclose(np.mean(results.fit_params['coefficients']), 0.005296895582675213, rel_tol=1e7)
        assert results.id == 'P66'

        scorer = Spectrum(80, False, False)
        results = scorer.fit(self.data)
        temp_dict = dict([(k, results.fit_params[k]) for k in results.fit_params.keys() if k != 'coefficients'])
        _dict_isclose(temp_dict, {'noise std': 0.16772880155008577,
                                  'order': 80,
                                  'peak': 22.28018448238728,
                                  'preprocessing': False,
                                  'period ci': None,
                                  'id': 'P66'}, rtol=1e-7)

        assert isclose(np.mean(results.fit_params['coefficients']), 0.0028957523385374885, rel_tol=1e7)
        assert isclose(np.mean(results.freq), 2.0833333333333333e-05, rel_tol=1e-7)
        assert isclose(np.mean(results.spectrum), 93.34486398942124, rel_tol=1e-7)

    def test_loading(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'ar.ada.long'))
        assert isinstance(loaded, ARResults)
        temp_dict = dict([(k, loaded.fit_params[k]) for k in loaded.fit_params.keys() if k != 'coefficients'])
        _dict_isclose(temp_dict, {'noise std': 0.16772880155008577,
                                  'order': 80,
                                  'peak': 22.28018448238728,
                                  'preprocessing': False,
                                  'period ci': None,
                                  'id': 'P66'}, rtol=1e-7)
        assert isclose(np.mean(loaded.fit_params['coefficients']), 0.0028957523385374885, rel_tol=1e7)
        assert isclose(np.mean(loaded.freq), 2.0833333333333333e-05, rel_tol=1e-7)
        assert isclose(np.mean(loaded.spectrum), 93.34486398942124, rel_tol=1e-7)
        assert loaded.id == 'P66'

    def test_exporting(self):
        loaded = FileManager.load_file(os.path.join(DIRECTORY, 'ar.ada.long'))
        loaded.export(os.path.join(DIRECTORY, 'temp_ar.ada.long'))
        exported = FileManager.load_file(os.path.join(DIRECTORY, 'temp_ar.ada.long'))
        assert isinstance(exported, ARResults)
        _dict_isclose(loaded.fit_params, exported.fit_params, rtol=1e-7)
        assert np.allclose(loaded.freq, exported.freq)
        assert np.allclose(loaded.spectrum, exported.spectrum)
        assert exported.id == loaded.id
        os.remove(os.path.join(DIRECTORY, 'temp_ar.ada.long'))
