import unittest
import os
import numpy as np
from glob import glob
import shutil

from bci_core import utils as ana_utils
from bci_core.online import model_loader
from training import train_model, model_saver
from dataloaders import neo
from online_sim import simulation
from validation import val_by_epochs


class TestOnlineSim(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        root_path = './tests/data'

        raw, cls.event_id = neo.raw_loader(root_path, {'flex': ['1', '2']})
        cls.raw = raw
        # split into 2 pieces
        t_min, t_max = raw.times[0], raw.times[-1]
        t_mid = raw.times[len(raw.times) // 2]
        raw_train = raw.copy().crop(tmin=t_min, tmax=t_mid, include_tmax=True)
        cls.raw_val = raw.copy().crop(tmin=t_mid, tmax=t_max)

        # reconstruct single event for validation
        if cls.raw_val.annotations.onset[0] > t_mid:
            # correct time by first timestamp
            cls.raw_val.annotations.onset -= t_mid
        
        # train with the first half
        model = train_model(raw_train, event_id=cls.event_id, model_type='baseline')
        model_saver(model, './tests/data/', 'baseline', 'test', cls.event_id)
        cls.model_path = glob(os.path.join('./tests/data/', 'test', '*.pkl'))[0]
    
    @classmethod
    def tearDownClass(cls) -> None:
        shutil.rmtree(os.path.join('./tests/data/', 'test'))
        return super().tearDownClass()
    
    def test_event_metric(self):
        event_gt = np.array([[0, 0, 0], [5, 0, 1], [7, 0, 0], [9, 0, 2]])
        event_pred = np.array([[1, 0, 0], [4, 0, 1], [6, 0, 1], [7, 0, 0], [10, 0, 1], [11, 0, 2]])
        fs = 1
        precision, recall, f1_score = ana_utils.event_metric(event_gt, event_pred, fs, ignore_event=(0,))
        self.assertEqual(f1_score, 2 / 3)
        self.assertEqual(precision, 1 / 2)
        self.assertEqual(recall, 1)

    def test_sim(self):
        model = model_loader(self.model_path, 
                             state_change_threshold=0.7,
                             state_trans_prob=0.7)
        metric_hmm, metric_nohmm, fig_pred = simulation(self.raw, self.event_id, model=model, epoch_length=1., step_length=0.1)
        fig_pred.savefig('./tests/data/pred.pdf')   
        self.assertTrue(metric_hmm[-2] > 0.3)  # f1-score (with hmm)
        self.assertTrue(metric_nohmm[-2] < 0.15)  # f1-score (without hmm)
    
    def test_val_model(self):
        metrices, fig_conf = val_by_epochs(self.raw_val, self.model_path, self.event_id, 1.)
        fig_conf.savefig('./tests/data/conf.pdf')
        self.assertGreater(metrices[0], 0.85)
        self.assertGreater(metrices[1], 0.7)
        self.assertGreater(metrices[2], 0.7)


if __name__ == '__main__':
    unittest.main()