# -*- coding: utf-8 -*-
"""Tests for reading data using io module."""

from ada.io.file_manager import FileManager
from ada.io.geneactiv import RawGeneActiv, GeneActivMVM
from ada.io.mesa import Mesa
from ada.io.actigraph import RawActiGraph
from ada.data_containers.generic import GenericData
from ada.io.acti_eeg import ActiPSG
from ada.data_containers.scored import PSGScore
from unittest import TestCase
import numpy as np
import json
import os
import pytest


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


@pytest.mark.order(0)
class TestReadingActi:
    def test_reading_GeneActiv(self):
        # Csv testing
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        assert isinstance(csv, GeneActivMVM)
        metadata = {'raw': csv.metadata, 'epoch': csv.epoching_method_metadata}
        with open(os.path.join(DIRECTORY, 'GeneActiveCsvMetadata.json'), 'r') as f:
            metadata = json.load(f)
        fs = 1 / 60
        data = np.load(os.path.join(DIRECTORY, 'GeneActiveCsvData.npy'))
        assert np.all(data == csv.data)
        assert fs == csv.fs
        assert csv.channel_names == ['x', 'y', 'z', 'lux', 'button', 'temperature', 'timestamp', 'mvm', 'x_std', 'y_std', 'z_std', 'peak_lux']
        assert csv.id == metadata['raw']['Subject Code']
        TestCase().assertDictEqual(metadata['raw'], csv.metadata)
        TestCase().assertDictEqual(metadata['epoch'], csv.epoching_method_metadata)
        # Binary testing
        binary = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBin.bin'))
        assert isinstance(binary, RawGeneActiv)
        assert binary.channel_names == ['x', 'y', 'z', 'lux', 'button', 'temperature', 'timestamp']
        fs = 100
        data = np.load(os.path.join(DIRECTORY, 'GeneActiveBinData.npy'))
        with open(os.path.join(DIRECTORY, 'GeneActiveBinMetadata.json'), 'r') as f:
            metadata = json.load(f)
        assert np.all(data[:5, :] == binary.data[:5, :])
        assert np.all(data[6, :] == binary.data[6, :])
        assert np.allclose(data[5, :], binary.data[5, :], 1e-2, 1e-1)  # temperature is a bit different
        assert fs == binary.fs
        assert binary.id == binary.metadata['Subject Code']
        TestCase().assertDictEqual(metadata, binary.metadata)
        assert binary.dynamic_range == 8.0
        assert binary.stationary_variance == 0.0004563


    def test_converting_to_generic(self):
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        generic = GenericData.from_nongeneric(csv)
        assert isinstance(generic, GenericData)
        assert csv.fs == generic.fs
        assert np.all(csv.data == generic.data)
        assert csv.channel_names == generic.channel_names
        assert generic.id == 'Ania'
        TestCase().assertDictEqual(csv.metadata, generic.metadata)
        TestCase().assertDictEqual(csv.epoching_method_metadata, generic.epoching_method_metadata)
        
        generic = GenericData.from_nongeneric(FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBin.bin')))
        assert generic.dynamic_range == 8.0
        assert generic.stationary_variance == 0.0004563
        assert generic.id == 'DK'


    def test_loading_generic(self):
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        generic = FileManager.load_file(os.path.join(DIRECTORY, 'genericExample.ada'))
        assert isinstance(generic, GenericData)
        assert csv.fs == generic.fs
        assert csv.channel_names == generic.channel_names
        assert np.all(csv.data == generic.data)
        assert generic.id == csv.id
        TestCase().assertDictEqual(csv.metadata, generic.metadata)
        TestCase().assertDictEqual(csv.epoching_method_metadata, generic.epoching_method_metadata)

    
    def test_exporting_GeneActiv(self):
        # From .csv to .csv
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        csv.export(os.path.join(DIRECTORY, 'GeneActiveCsvExported.csv'))
        csv_exported = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsvExported.csv'))
        assert np.allclose(csv.data, csv_exported.data)
        # Metadata partially changes (e.g. software name), leave it for now
        # TestCase().assertDictEqual(csv.metadata, csv_exported.metadata)
        TestCase().assertDictEqual(csv.epoching_method_metadata, csv_exported.epoching_method_metadata)
        # From .bin to .csv"
        binary = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBin.bin'))
        binary.export(os.path.join(DIRECTORY, 'GeneActiveBinExported.csv'))
        csv_exported = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveBinExported.csv'))
        assert np.allclose(binary.data, csv_exported.data)
        assert binary.id == csv_exported.id
        # Afer exporting to .csv metadata changes to be coherent with manufacturers format, so no testing for that.
        os.remove(os.path.join(DIRECTORY, 'GeneActiveCsvExported.csv'))
        os.remove(os.path.join(DIRECTORY, 'GeneActiveBinExported.csv'))

    
    def test_exporting_generic(self):
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'GeneActiveCsv.csv'))
        generic = GenericData.from_nongeneric(csv)
        generic.export(os.path.join(DIRECTORY, 'genericExported.ada'))
        generic_loaded = FileManager.load_file(os.path.join(DIRECTORY, 'genericExported.ada'))
        assert isinstance(generic_loaded, GenericData)
        assert generic.channel_names == generic_loaded.channel_names
        assert np.allclose(generic.data, generic_loaded.data)
        assert generic.fs == generic_loaded.fs
        assert generic.id == generic_loaded.id
        TestCase().assertDictEqual(generic.metadata, generic_loaded.metadata)
        os.remove(os.path.join(DIRECTORY, 'genericExported.ada'))

    
    def test_reading_mesa(self):
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'MESA.csv'))
        assert isinstance(csv, Mesa)
        assert csv.fs == 1 / 30
        assert csv.channel_names == ['score',
                                     'timestamp',
                                     'offwrist',
                                     'activity',
                                     'marker',
                                     'whitelight',
                                     'redlight',
                                     'greenlight',
                                     'bluelight',
                                     'interval',
                                     'dayofweek',
                                     'daybymidnight',
                                     'daybynoon']
        data = np.array([[1.000e+00, 0.000e+00, 1.000e+00, 1.000e+00],
                         [0.000e+00, 3.000e+01, 6.000e+01, 9.000e+01],
                         [0.000e+00, 0.000e+00, 0.000e+00, 0.000e+00],
                         [1.070e+02, 0.000e+00, np.nan, 3.200e+01],
                         [0.000e+00, 0.000e+00, np.nan, 0.000e+00],
                         [7.760e+00, 1.358e+01, np.nan, 1.750e+00],
                         [1.750e+00, 2.700e+00, np.nan, 7.200e-01],
                         [4.370e-01, 8.560e-01, np.nan, 1.730e-02],
                         [5.960e-02, 1.140e-01, np.nan, 5.900e-03],
                         [0.000e+00, 5.000e-01, np.nan, 1.000e+00],
                         [5.000e+00, 5.000e+00, 5.000e+00, 5.000e+00],
                         [1.000e+00, 1.000e+00, 1.000e+00, 1.000e+00],
                         [1.000e+00, 1.000e+00, 1.000e+00, 1.000e+00]])
        assert np.allclose(csv.data, data, equal_nan=True)
        assert csv.id == '1'
        TestCase().assertDictEqual(csv.metadata, {'first_line': 378, 'mesa_id': 1, 'sampling_frequency': 0.03333333333333333, 'start_hour': '18:57:30'})
        TestCase().assertDictEqual(csv.epoching_method_metadata, {'epoching method': 'MESA', 'to_score channel': 3})
        TestCase().assertDictEqual(csv.scoring_method_metadata, {'algorithm': 'MESA'})


    def test_exporting_mesa(self):
        csv = FileManager.load_file(os.path.join(DIRECTORY, 'MESA.csv'))
        csv.export(os.path.join(DIRECTORY, 'temp_mesa.csv'))
        csv_exported = FileManager.load_file(os.path.join(DIRECTORY, 'temp_mesa.csv'))
        os.remove(os.path.join(DIRECTORY, 'temp_mesa.csv'))
        print(csv.data)
        print(csv_exported.data)
        assert isinstance(csv_exported, Mesa)
        assert csv.fs == csv_exported.fs
        assert csv.channel_names == csv_exported.channel_names
        assert np.allclose(csv.data, csv_exported.data, equal_nan=True)
        assert csv.id == csv_exported.id
        TestCase().assertDictEqual(csv.metadata, csv_exported.metadata)
        TestCase().assertDictEqual(csv.epoching_method_metadata, csv_exported.epoching_method_metadata)
        TestCase().assertDictEqual(csv.scoring_method_metadata, csv_exported.scoring_method_metadata)


    def test_reading_actigraph_gt3x(self):
        # File neo.gt3x from https://github.com/actigraph/pygt3x/tree/main, thanks
        # Btw for some reason timestamp here is in 1990, but code is ok (according to gt3x specification it IS unix timestamp)
        data = FileManager.load_file(os.path.join(DIRECTORY, 'neo.gt3x'))
        assert isinstance(data, RawActiGraph)
        assert data.fs == 30
        assert data.metadata == {'acceleration_max': 0.0,
                                 'acceleration_min': 0.0,
                                 'acceleration_scale': 0.0,
                                 'age': None, 'battery_voltage': 4.26,
                                 'board_revision': '2',
                                 'device_type': 'GT3XPlus',
                                 'dominance': None,
                                 'download_date': 637755989984676604,
                                 'firmware': '2.5.0',
                                 'height': None,
                                 'last_sample_time': 0,
                                 'limb': None,
                                 'mass': None,
                                 'race': None,
                                 'sample_rate': 30,
                                 'serial_number':
                                 'NEO1C04110003',
                                 'sex': None,
                                 'side': None,
                                 'start_date': 637755981000000000,
                                 'stop_date': 637755987000000000,
                                 'subject_name': 'NEO1C04110003',
                                 'timezone': '-06:00:00',
                                 'unexpected_resets': '0',
                                 'first_sample': 637755981.0}
        assert data.channel_names == ['timestamp', 'x', 'y', 'z']
        assert data.first_sample_timestamp == 637755981.0
        assert data.id == 'NEO1C04110003'
        assert np.allclose(np.mean(data.data, axis=1), [3.00050000e+02, -6.40106096e-02, 6.23935787e-02, -6.77247327e-01])

    
    def test_reading_actigraph_csv(self):
        data = FileManager.load_file(os.path.join(DIRECTORY, 'Actigraph.csv'))
        assert isinstance(data, RawActiGraph)
        assert data.fs == 30
        assert data.metadata == {'sample_rate': 30.0,
                                 'firmware': '1.9.2',
                                 'serial_number':
                                 'MOS2D36170057',
                                 'subject_name': '',
                                 'device_type': 'wGT3X-BT',
                                 'battery_voltage': 3.91,
                                 'start_date': 1638097200.0,
                                 'stop_date': 1638097200.034,
                                 'download_date': 1639573190.0,
                                 'first_sample': 1638097200.0}
        assert data.first_sample_timestamp == 1638097200.0
        assert np.allclose(data.data, [[0., 0.03399992], [0.74, 0.63], [-0.52, -0.543], [-0.836, -0.84]])
        assert data.channel_names == ['timestamp', 'x', 'y', 'z']
        assert data.id == ''


    def test_exporting_acigraph(self):
        data = FileManager.load_file(os.path.join(DIRECTORY, 'Actigraph.csv'))
        data.export(os.path.join(DIRECTORY, 'test_actigraph.csv'))
        exported = FileManager.load_file(os.path.join(DIRECTORY, 'test_actigraph.csv'))
        assert exported.fs == data.fs
        assert np.allclose(exported.data, data.data)
        assert data.channel_names == exported.channel_names
        assert data.first_sample_timestamp == exported.first_sample_timestamp
        TestCase().assertDictEqual(exported.metadata, data.metadata)
        os.remove(os.path.join(DIRECTORY, 'test_actigraph.csv'))


@pytest.mark.order(1)
class TestReadingPsg:
    def test_stages_heuristic(self):
        tags = [dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['w', 1., 21.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['rem', 21., 41.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['s1', 41., 61.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['s2', 61., 81.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['stadium 3', 81., 101.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['szybka koza', 101., 121.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['artifacts', 121., 141.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['unin/art', 141., 161.]))]
        tags = [e['name'] for e in tags]
        collapsed_stages = np.array([1, 0, 0, 0, 0, 1, None, None], dtype=float)
        psg_score = PSGScore(tags, 1., 141., 20, True)
        assert np.all(psg_score.psg_stages[:-2] == collapsed_stages[:-2])
        assert np.all(np.isnan(psg_score.psg_stages[-2:]))


    def test_custom_heuristic(self):
        tags = [dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['w', 1., 21.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['rem', 21., 41.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['s1', 41., 61.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['s2', 61., 81.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['stadium 3', 81., 101.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['szybka koza', 101., 121.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['artifacts', 121., 141.])),
                dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['unin/art', 141., 161.]))]
        tags = [e['name'] for e in tags]
        
        def heuristic(tag):
            if 'koza' in tag:
                return 1
            return 0
        
        collapsed_stages = np.zeros(8, dtype=float)
        collapsed_stages[-3] = 1.
        psg_score = PSGScore(tags, 1., 141., 20, heuristic)
        assert np.all(collapsed_stages == psg_score.psg_stages)

    
    def test_constructing_trimmed(self):
        raw = FileManager.load_file(os.path.join(DIRECTORY, 'epochingExample.ada'))
        tags = [dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['w', 1706472231.5 + 100 + i * 20, 1706472231.5 + 100 + (i + 1) * 20])) for i in range(1200 // 20)]
        acti_psg = ActiPSG.from_continous_stages(raw, tags, 1706472231.5 + 100, False, True)
        assert isinstance(acti_psg, ActiPSG)
        assert acti_psg.acti_data.fs == 75.
        assert np.all(raw.data[:, 100 * 75:(100 + 1200) * 75] == acti_psg.acti_data.data)
        assert len(acti_psg.acti_data.timestamp) == len(acti_psg.eeg_data.psg_stages)
        assert np.isclose(acti_psg.first_sample_timestamp, 1706472231.5 + 100)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, acti_psg.eeg_data.start_timestamp)
        assert np.isclose(acti_psg.acti_data.last_sample_timestamp, acti_psg.eeg_data.end_timestamp)
        TestCase().assertDictEqual(raw.metadata, acti_psg.acti_data.metadata)


    def test_constructing_extended(self):
        raw = FileManager.load_file(os.path.join(DIRECTORY, 'epochingExample.ada'))
        tags = [dict(zip(['name', 'start_timestamp', 'end_timestamp'], ['w', 1706472231.5 + 100 + i * 20, 1706472231.5 + 100 + (i + 1) * 20])) for i in range(1200 // 20)]
        acti_psg = ActiPSG.from_continous_stages(raw, tags, 1706472231.5 + 100, False, False)
        assert isinstance(acti_psg, ActiPSG)
        assert acti_psg.acti_data.fs == 75.
        assert np.all(raw.data == acti_psg.acti_data.data)
        assert len(acti_psg.acti_data.timestamp) == len(acti_psg.eeg_data.psg_stages)
        assert np.isclose(acti_psg.first_sample_timestamp, 1706472231.5 + 100)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, acti_psg.eeg_data.start_timestamp)
        assert np.isclose(acti_psg.acti_data.last_sample_timestamp, acti_psg.eeg_data.end_timestamp)
        TestCase().assertDictEqual(raw.metadata, acti_psg.acti_data.metadata)
    
    
    def test_loading_from_tags(self):
        acti_psg = FileManager.load_ActiPSG_from_tags(os.path.join(DIRECTORY, 'ActiPSG.ada'), os.path.join(DIRECTORY, 'ActiPSG.obci.tag'), os.path.join(DIRECTORY, 'ActiPSG.obci.xml'))
        acti = FileManager.load_file(os.path.join(DIRECTORY, 'ActiPSG.ada'))
        assert isinstance(acti_psg, ActiPSG)
        assert acti.fs == acti_psg.acti_data.fs
        assert np.all(acti.data == acti_psg.acti_data.data)
        assert np.all(np.array([0, 0, 1, 1, 0, 0, 1], dtype=float) == acti_psg.eeg_data.psg_stages[:-1])
        assert np.isnan(acti_psg.eeg_data.psg_stages[-1])
        TestCase().assertDictEqual(acti.metadata, acti_psg.acti_data.metadata)
        TestCase().assertDictEqual(acti.epoching_method_metadata, acti_psg.acti_data.epoching_method_metadata)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, 1710204506.187291)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, acti_psg.eeg_data.start_timestamp)
        assert np.isclose(acti_psg.acti_data.last_sample_timestamp, acti_psg.eeg_data.end_timestamp)


    def test_loading(self):
        acti_psg = FileManager.load_file(os.path.join(DIRECTORY, 'ActiPSG.zip'))
        acti = FileManager.load_file(os.path.join(DIRECTORY, 'ActiPSG.ada'))
        assert isinstance(acti_psg, ActiPSG)
        assert acti.fs == acti_psg.acti_data.fs
        assert np.all(acti.data == acti_psg.acti_data.data)
        assert np.all(np.array([0, 0, 1, 1, 0, 0, 1], dtype=float) == acti_psg.eeg_data.psg_stages[:-1])
        assert np.isnan(acti_psg.eeg_data.psg_stages[-1])
        TestCase().assertDictEqual(acti.metadata, acti_psg.acti_data.metadata)
        TestCase().assertDictEqual(acti.epoching_method_metadata, acti_psg.acti_data.epoching_method_metadata)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, 1710204506.187291)
        assert np.isclose(acti_psg.acti_data.first_sample_timestamp, acti_psg.eeg_data.start_timestamp)
        assert np.isclose(acti_psg.acti_data.last_sample_timestamp, acti_psg.eeg_data.end_timestamp)


    def test_exporting(self):
        acti_psg = FileManager.load_file(os.path.join(DIRECTORY, 'ActiPSG.zip'))
        out_path = os.path.join(DIRECTORY, 'temp_psg.zip')
        acti_psg.export(out_path)
        acti_psg_2 = FileManager.load_file(out_path)
        os.remove(out_path)
        assert isinstance(acti_psg_2, ActiPSG)
        assert acti_psg.acti_data.fs == acti_psg_2.acti_data.fs
        assert np.all(acti_psg.acti_data.data == acti_psg_2.acti_data.data)
        assert np.all(acti_psg.eeg_data.psg_stages[:-1] == acti_psg_2.eeg_data.psg_stages[:-1])
        assert np.isnan(acti_psg_2.eeg_data.psg_stages[-1])
        TestCase().assertDictEqual(acti_psg.acti_data.metadata, acti_psg_2.acti_data.metadata)
        TestCase().assertDictEqual(acti_psg.acti_data.epoching_method_metadata, acti_psg_2.acti_data.epoching_method_metadata)
        assert acti_psg.eeg_data.start_timestamp == acti_psg_2.eeg_data.start_timestamp
        assert acti_psg.eeg_data.end_timestamp == acti_psg_2.eeg_data.end_timestamp
