"""Tests for the ibllib.qc.task_qc_viewer package."""
import os
import unittest
from unittest import mock

from one.api import ONE
import numpy as np

from ibllib.pipes.ephys_preprocessing import EphysTrials
from ibllib.pipes.training_preprocessing import TrainingTrials
from ibllib.pipes.behavior_tasks import HabituationTrialsBpod, ChoiceWorldTrialsNidq, ChoiceWorldTrialsBpod, PassiveTask
from ibllib.qc.task_qc_viewer.task_qc import get_bpod_trials_task, show_session_task_qc, QcFrame
from ibllib.qc.task_metrics import TaskQC
from ibllib.tests import TEST_DB


MOCK_QT = os.environ.get('IBL_MOCK_QT', True)
"""bool: If true, do not run the QT application."""


class TestTaskQC(unittest.TestCase):
    """Tests for ibllib.qc.task_qc_viewer.task_qc module."""

    def setUp(self):
        self.one = ONE(**TEST_DB, mode='local')
        """Some testing environments do not have the correct QT libraries. It is difficult to
        ensure Qt is installed correctly as Anaconda, OpenCV, and system QT installations can
        disrupt the lib paths. If MOCK_QT is true, the QC application is never run."""
        if MOCK_QT:
            qt_mock = mock.patch('ibllib.qc.task_qc_viewer.ViewEphysQC.viewqc')
            qt_mock.start()
            self.addCleanup(qt_mock.stop)

    def test_get_bpod_trials_task(self):
        """Test get_bpod_trials_task function."""
        task = TrainingTrials('foo/bar', one=self.one)
        bpod_task = get_bpod_trials_task(task)
        self.assertIs(task, bpod_task)

        task = HabituationTrialsBpod('foo/bar', one=self.one,
                                     protocol_number=0, protocol='habituationChoiceWorld', collection='raw_task_data_00')
        bpod_task = get_bpod_trials_task(task)
        self.assertIs(task, bpod_task)

        task = ChoiceWorldTrialsNidq('foo/bar', one=self.one,
                                     protocol_number=2, protocol='ephysChoiceWorld', collection='raw_task_data_02')
        bpod_task = get_bpod_trials_task(task)
        self.assertIs(bpod_task.__class__, ChoiceWorldTrialsBpod)
        self.assertEqual(bpod_task.protocol_number, 2)
        self.assertEqual(bpod_task.protocol, 'ephysChoiceWorld')
        self.assertEqual(bpod_task.collection, 'raw_task_data_02')
        self.assertIs(bpod_task.one, self.one)

        task = EphysTrials('foo/bar', one=self.one)
        bpod_task = get_bpod_trials_task(task)
        self.assertIsInstance(bpod_task, TrainingTrials)

    @mock.patch('ibllib.qc.task_qc_viewer.task_qc.qt.run_app')
    @mock.patch('ibllib.qc.task_qc_viewer.task_qc.get_trials_tasks')
    def test_show_session_task_qc(self, trials_tasks_mock, run_app_mock):
        """Test show_session_task_qc function."""
        trials_tasks_mock.return_value = []
        session_path = 'foo/bar/subject/2023-01-01/001'
        self.assertRaises(ValueError, show_session_task_qc, session_path, one=self.one)
        self.assertRaises(TypeError, show_session_task_qc, session_path, one=self.one, protocol_number=-2)
        self.assertRaises(ValueError, show_session_task_qc, session_path, one=self.one, protocol_number=1)

        passive_task = PassiveTask('foo/bar', protocol='_iblrig_passiveChoiceWorld', protocol_number=0)
        trials_tasks_mock.return_value = [passive_task]
        self.assertRaises(ValueError, show_session_task_qc, session_path, one=self.one, protocol_number=0)
        self.assertRaises(ValueError, show_session_task_qc, session_path, one=self.one)

        # Set up QC mock
        qc_mock = mock.Mock(spec=TaskQC, unsafe=True)
        qc_mock.metrics = {'foo': .7}
        qc_mock.compute_session_status.return_value = ('Fail', qc_mock.metrics, {'foo': 'FAIL'})
        qc_mock.extractor.data = {'intervals': np.array([[0, 1]])}
        qc_mock.extractor.frame_ttls = qc_mock.extractor.audio_ttls = qc_mock.extractor.bpod_ttls = mock.MagicMock()

        active_task = mock.Mock(spec=ChoiceWorldTrialsNidq, unsafe=True)
        active_task.run_qc.return_value = qc_mock
        active_task.name = 'Trials_activeChoiceWorld_01'
        trials_tasks_mock.return_value = [passive_task, active_task]
        qc = show_session_task_qc(session_path, one=self.one)

        self.assertIsInstance(qc, QcFrame)
        self.assertIsInstance(qc.qc, TaskQC)
        self.assertCountEqual(qc.get_wheel_data(), ('re_ts', 're_pos'))
        active_task.run_qc.assert_called_once_with(update=False)
        self.assertEqual('remote', active_task.location)
        active_task.setUp.assert_called_once()
        active_task.assert_expected_inputs.assert_not_called()
        run_app_mock.assert_called_once()

        active_task.reset_mock(return_value=False)
        with mock.patch('ibllib.qc.task_qc_viewer.task_qc.get_bpod_trials_task', return_value=active_task) as \
                get_bpod_trials_task_mock:
            show_session_task_qc(session_path, one=self.one, local=True, bpod_only=True)
            # Should be called in bpod_only mode
            get_bpod_trials_task_mock.assert_called_once_with(active_task)
        # Should be called in local mode
        active_task.assert_expected_inputs.assert_called_once_with(raise_error=True)

        # If QcFrame instance passed, should use this and return it
        self.assertIs(show_session_task_qc(qc, one=self.one), qc)
        # If passing TaskQC object, should not call trials_tasks_mock
        trials_tasks_mock.reset_mock()
        show_session_task_qc(qc_mock, one=self.one)
        self.assertIsInstance(qc, QcFrame)
        trials_tasks_mock.assert_not_called()


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