from .test_utils import PySparkTestCase
from sparkts.datetimeindex import *
import pandas as pd

class DateTimeIndexTestCase(PySparkTestCase):
    def test_frequencies(self):
        bd = BusinessDayFrequency(1, 1, self.sc)
        self.assertEqual(bd.days(), 1)
        
        hf = HourFrequency(4, self.sc)
        self.assertEqual(hf.hours(), 4)
    
    def test_uniform(self):
        freq = DayFrequency(3, self.sc)
        self.assertEqual(freq.days(), 3)
        start = '2015-04-10'
        index = uniform(start, periods=5, freq=freq, sc=self.sc)
        index2 = uniform(start, end='2015-04-22', freq=freq, sc=self.sc)
        self.assertEqual(index, index2)
        
        self.assertEqual(len(index), 5)
        self.assertEqual(index.first(), pd.to_datetime('2015-04-10'))
        self.assertEqual(index.last(), pd.to_datetime('2015-04-22'))
        subbydate = index[pd.to_datetime('2015-04-13'):pd.to_datetime('2015-04-19')]
        subbyloc = index.islice(1, 4)
        self.assertEqual(subbydate, subbyloc)
        self.assertEqual(subbydate.first(), pd.to_datetime('2015-04-13'))
        self.assertEqual(subbydate.last(), pd.to_datetime('2015-04-19'))
        self.assertEqual(subbydate.datetime_at_loc(0), pd.to_datetime('2015-04-13'))
        self.assertEqual(subbydate[pd.to_datetime('2015-04-13')], 0)

    def test_irregular(self):
        pd_index = pd.date_range('2015-04-10', periods=5, freq='3D')
        dt_index = irregular(pd_index, self.sc)

        self.assertEqual(len(dt_index), 5)
        self.assertEqual(dt_index.first(), pd.to_datetime('2015-04-10'))
        self.assertEqual(dt_index.last(), pd.to_datetime('2015-04-22'))
        subbydate = dt_index[pd.to_datetime('2015-04-13'):pd.to_datetime('2015-04-19')]
        subbyloc = dt_index.islice(1, 4)
        self.assertEqual(subbydate, subbyloc)
        self.assertEqual(subbydate.first(), pd.to_datetime('2015-04-13'))
        self.assertEqual(subbydate.last(), pd.to_datetime('2015-04-19'))
        self.assertEqual(subbydate.datetime_at_loc(0), pd.to_datetime('2015-04-13'))
        self.assertEqual(subbydate[pd.to_datetime('2015-04-13')], 0)

        pd_index2 = dt_index.to_pandas_index()
        self.assertTrue(pd_index.equals(pd_index2), str(pd_index) + "!=" + str(pd_index2))


[((33, 19, 33, 68), 'pandas.date_range', 'pd.date_range', (), '', True, 'import pandas as pd\n'), ((22, 40, 22, 68), 'pandas.to_datetime', 'pd.to_datetime', ({(22, 55, 22, 67): '"""2015-04-10"""'}, {}), "('2015-04-10')", True, 'import pandas as pd\n'), ((23, 39, 23, 67), 'pandas.to_datetime', 'pd.to_datetime', ({(23, 54, 23, 66): '"""2015-04-22"""'}, {}), "('2015-04-22')", True, 'import pandas as pd\n'), ((27, 44, 27, 72), 'pandas.to_datetime', 'pd.to_datetime', ({(27, 59, 27, 71): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((28, 43, 28, 71), 'pandas.to_datetime', 'pd.to_datetime', ({(28, 58, 28, 70): '"""2015-04-19"""'}, {}), "('2015-04-19')", True, 'import pandas as pd\n'), ((29, 55, 29, 83), 'pandas.to_datetime', 'pd.to_datetime', ({(29, 70, 29, 82): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((37, 43, 37, 71), 'pandas.to_datetime', 'pd.to_datetime', ({(37, 58, 37, 70): '"""2015-04-10"""'}, {}), "('2015-04-10')", True, 'import pandas as pd\n'), ((38, 42, 38, 70), 'pandas.to_datetime', 'pd.to_datetime', ({(38, 57, 38, 69): '"""2015-04-22"""'}, {}), "('2015-04-22')", True, 'import pandas as pd\n'), ((42, 44, 42, 72), 'pandas.to_datetime', 'pd.to_datetime', ({(42, 59, 42, 71): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((43, 43, 43, 71), 'pandas.to_datetime', 'pd.to_datetime', ({(43, 58, 43, 70): '"""2015-04-19"""'}, {}), "('2015-04-19')", True, 'import pandas as pd\n'), ((44, 55, 44, 83), 'pandas.to_datetime', 'pd.to_datetime', ({(44, 70, 44, 82): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((24, 26, 24, 54), 'pandas.to_datetime', 'pd.to_datetime', ({(24, 41, 24, 53): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((24, 55, 24, 83), 'pandas.to_datetime', 'pd.to_datetime', ({(24, 70, 24, 82): '"""2015-04-19"""'}, {}), "('2015-04-19')", True, 'import pandas as pd\n'), ((39, 29, 39, 57), 'pandas.to_datetime', 'pd.to_datetime', ({(39, 44, 39, 56): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((39, 58, 39, 86), 'pandas.to_datetime', 'pd.to_datetime', ({(39, 73, 39, 85): '"""2015-04-19"""'}, {}), "('2015-04-19')", True, 'import pandas as pd\n'), ((30, 35, 30, 63), 'pandas.to_datetime', 'pd.to_datetime', ({(30, 50, 30, 62): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n'), ((45, 35, 45, 63), 'pandas.to_datetime', 'pd.to_datetime', ({(45, 50, 45, 62): '"""2015-04-13"""'}, {}), "('2015-04-13')", True, 'import pandas as pd\n')]