import unittest

from riptable import *


class PGroupby_Test(unittest.TestCase):
    def test_groupbyhash_cutoffs(self):
        arr = np.random.randint(0, 5, 15)
        arr2 = np.random.randint(0, 5, 15)
        arr3 = hstack([arr, arr2])

        gb_arr = groupbyhash([arr])
        gb_arr2 = groupbyhash([arr2])
        gb_arr3 = groupbyhash([arr3], cutoffs=FA([15, 30], dtype=np.int64))

        ikey = gb_arr3['iKey']
        ifirst_vals = gb_arr3['iFirstKey'][0]
        ifirst_cutoffs = gb_arr3['iFirstKey'][1]

        self.assertTrue(bool(np.all(gb_arr['iKey'] == ikey[:15])))
        self.assertTrue(
            bool(np.all(gb_arr['iFirstKey'] == ifirst_vals[: ifirst_cutoffs[0]]))
        )

        self.assertTrue(bool(np.all(gb_arr2['iKey'] == ikey[15:])))
        self.assertTrue(
            bool(np.all(gb_arr2['iFirstKey'] == ifirst_vals[ifirst_cutoffs[0] :]))
        )
        self.assertEqual(
            gb_arr3['unique_count'], (gb_arr['unique_count'] + gb_arr2['unique_count'])
        )

    def test_cutoffs_vs_multikey(self):
        arr = np.random.randint(0, 5, 15)
        arr2 = np.random.randint(0, 5, 15)
        arr3 = hstack([arr, arr2])
        mk_col = hstack([zeros(15), ones(15)])

        gb_mk = groupbyhash([mk_col, arr3])
        gb_arr3 = groupbyhash([arr3], cutoffs=FA([15, 30], dtype=np.int64))
        ikey = gb_arr3['iKey']
        ifirst_vals = gb_arr3['iFirstKey'][0]
        ifirst_cutoffs = gb_arr3['iFirstKey'][1]

        self.assertTrue(
            bool(
                np.all(
                    gb_mk['iFirstKey'][: ifirst_cutoffs[0]]
                    == ifirst_vals[: ifirst_cutoffs[0]]
                )
            )
        )
        self.assertTrue(
            bool(
                np.all(
                    gb_mk['iFirstKey'][ifirst_cutoffs[0] :]
                    == ifirst_vals[ifirst_cutoffs[0] :] + 15
                )
            )
        )
        self.assertEqual(gb_mk['unique_count'], gb_arr3['unique_count'])

    def test_single_unique_per_group(self):
        mk_col = hstack([zeros(15), ones(15)])
        pgb = groupbyhash([mk_col], cutoffs=FA([15, 30], dtype=np.int64))
        self.assertTrue(bool(np.all(pgb['iKey'] == 1)))
        self.assertTrue(bool(np.all(pgb['iFirstKey'][0] == 0)))
        self.assertEqual(pgb['unique_count'], 2)

        gb1 = groupbyhash([mk_col[:15]])
        gb2 = groupbyhash([mk_col[15:]])
        self.assertEqual(gb1['unique_count'], 1)
        self.assertEqual(gb2['unique_count'], 1)

    def test_single_cutoff(self):
        arr = np.random.randint(0, 10, 15)
        gb = groupbyhash([arr])
        pgb = groupbyhash([arr], cutoffs=FA([15], dtype=np.int64))

        self.assertTrue(bool(np.all(gb['iKey'] == pgb['iKey'])))
        self.assertTrue(bool(np.all(gb['iFirstKey'] == pgb['iFirstKey'][0])))
        self.assertEqual(len(pgb['iFirstKey'][1]), 1)
        self.assertEqual(pgb['iFirstKey'][1][0], pgb['unique_count'])

    def test_single_item_per_group(self):
        cutoffs = arange(10, dtype=np.int64) + 1
        gb = groupbyhash([arange(10)])
        pgb = groupbyhash([ones(10)], cutoffs=cutoffs)

        self.assertTrue(bool(np.all(gb['iKey'] == pgb['iFirstKey'][1])))
        self.assertEqual(gb['unique_count'], pgb['unique_count'])
        self.assertTrue(bool(np.all(pgb['iKey'] == 1)))
        self.assertTrue(bool(np.all(pgb['iFirstKey'][0] == 0)))

    def test_empty_cutoffs(self):
        a = arange(10)
        cutoffs = a.astype(np.int64) + 1
        f = logical(a % 2)
        no_change = FastArray([0, 1, 1, 2, 2, 3, 3, 4, 4, 5], dtype=np.int64)

        pgb = groupbyhash([ones(10)], filter=f, cutoffs=cutoffs)
        self.assertTrue(bool(np.all(pgb['iKey'] == f)))
        self.assertTrue(bool(np.all(pgb['iFirstKey'][1] == no_change)))


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