import unittest
from datetime import datetime

import numpy as np
from dateutil.tz import tzlocal
from nwbwidgets.utils.timeseries import (
    get_timeseries_tt, get_timeseries_maxt, get_timeseries_mint, get_timeseries_in_units, timeseries_time_to_ind,
    align_by_times, align_by_trials, align_by_time_intervals
)
from pynwb import NWBFile
from pynwb import TimeSeries
from pynwb.epoch import TimeIntervals


def test_get_timeseries_tt():
    data = list(range(100, 200, 10))
    ts = TimeSeries(name='test_timeseries', data=data, unit='m', starting_time=0.0, rate=1.0)

    tt = get_timeseries_tt(ts)
    np.testing.assert_array_equal(tt, [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])


def test_get_timeseries_tt_infstarting_time():
    data = list(range(100, 200, 10))
    ts = TimeSeries(name='test_timeseries', data=data, unit='m', starting_time=np.inf, rate=1.0)

    tt = get_timeseries_tt(ts)
    np.testing.assert_array_equal(tt, [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])


def test_get_timeseries_tt_negativeistop():
    data = list(range(100, 200, 10))
    ts = TimeSeries(name='test_timeseries', data=data, unit='m', starting_time=0., rate=1.0)

    tt = get_timeseries_tt(ts, istop=-1)
    np.testing.assert_array_equal(tt, [0., 1., 2., 3., 4., 5., 6., 7.])


def test_get_timeseries_in_units():
    data = list(range(100, 200, 10))
    timestamps = list(range(10))
    ts = TimeSeries(name='test_timeseries', data=data, unit='m', timestamps=timestamps, conversion=np.inf)
    data, unit = get_timeseries_in_units(ts)
    assert (unit is None)
    assert (data == [100, 110, 120, 130, 140, 150, 160, 170, 180, 190])


def test_align_by_trials():
    start_time = datetime(2017, 4, 3, 11, tzinfo=tzlocal())
    create_date = datetime(2017, 4, 15, 12, tzinfo=tzlocal())

    nwbfile = NWBFile(session_description='NWBFile for PSTH',
                      identifier='NWB123',
                      session_start_time=start_time,
                      file_create_date=create_date)

    data = list(range(100, 200, 10))
    timestamps = list(range(10))
    ts = TimeSeries(name='test_timeseries', data=data, unit='m', timestamps=timestamps)
    nwbfile.add_acquisition(ts)

    nwbfile.add_trial_column(name='stim', description='the visual stimuli during the trial')
    nwbfile.add_trial(start_time=0.0, stop_time=2.0, stim='person')
    nwbfile.add_trial(start_time=3.0, stop_time=5.0, stim='ocean')
    nwbfile.add_trial(start_time=6.0, stop_time=8.0, stim='desert')

    np.testing.assert_array_equal(align_by_trials(ts), np.array([[110], [140], [170]]))


class TimeSeriesTimeStampTestCase(unittest.TestCase):

    def setUp(self):
        data = list(range(100, 200, 10))
        timestamps = list(range(10))
        self.ts = TimeSeries(name='test_timeseries', data=data, unit='m', timestamps=timestamps)

    def test_get_timeseries_maxt(self):
        maxt = get_timeseries_maxt(self.ts)
        assert (maxt == 9)

    def test_get_timeseries_mint(self):
        mint = get_timeseries_mint(self.ts)
        assert (mint == 0)

    def test_timeseries_time_to_ind(self):
        assert (timeseries_time_to_ind(self.ts, 3) == 4)

    def test_align_by_times(self):
        assert (np.array_equal(align_by_times(self.ts, [0, 1, 2], [4, 5, 6]), np.array([[110, 120, 130, 140],
                                                                                        [120, 130, 140, 150],
                                                                                        [130, 140, 150, 160]])))

    def test_get_timeseries_tt_timestamp(self):
        tt = get_timeseries_tt(self.ts)
        np.testing.assert_array_equal(tt, [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])

    def test_align_by_time_intervals(self):
        intervals = TimeIntervals(name='Time Intervals')
        np.testing.assert_array_equal(align_by_time_intervals(
            timeseries=self.ts, intervals=intervals, stop_label=None), np.array([]))
