import os
import unittest
from math import isclose

from riptable import *
from riptable import Categorical, Dataset, FastArray, GroupBy, arange, isnan, isnotnan
from riptable.rt_enum import INVALID_DICT
from riptable.tests.utils import LikertDecision
from riptable.Utils.pandas_utils import dataset_as_pandas_df

num_list = [3, 4, 5, 1, 2]
num_list_sorted = [1, 2, 3, 4, 5]
arr_types = [
    np.int8,
    np.uint8,
    np.int16,
    np.uint16,
    np.int32,
    np.uint32,
    np.int64,
    np.uint64,
    np.float32,
    np.float64,
]
arr_types_string = [np.bytes_, np.str_]
test_data = {"bool": np.array([True, False, True, False, True], dtype=bool)}
for dt in arr_types + arr_types_string:
    test_data[dt.__name__] = np.array(num_list, dtype=dt)
test_data["categorical"] = Categorical([str(i) for i in num_list])
all_headers = list(test_data.keys())
ds = Dataset(test_data)
gb_funcs = ["sum", "mean", "first", "last", "median", "min", "max", "var"]
gb_nan_funcs = ["nansum", "nanmean", "nanmedian", "nanvar"]  #'rolling', 'cumsum', 'nth'


class Groupby_Test(unittest.TestCase):
    def test_math_ops_same_return(self):
        result_dict = {
            "sum": [5, 10],
            "nansum": [5, 10],
            # TODO: add support for min / max on strings
            "min": [1, 2],
            "max": [4, 5],
        }
        gb = ds.gb("bool")
        for name, correct in result_dict.items():
            call = getattr(gb, name)
            ds_result = call()
            for dt in arr_types:
                column = getattr(ds_result, dt.__name__)
                is_correct = bool(np.all(column == np.array(correct, dtype=dt)))
                self.assertTrue(
                    is_correct,
                    msg=f"Incorrect result for keys {ds_result.label_get_names()} in column with datatype {dt} after {name} operation.",
                )

            # make sure string columns are not included
            for dt in arr_types_string:
                was_included = hasattr(ds_result, dt.__name__)
                self.assertFalse(
                    was_included,
                    msg=f"String-like column {dt.__name__} was included in {name} math operation.",
                )

    def test_filter(self):
        correct = [1, 2]
        result = ds.gb("bool", filter=[False, False, False, True, True]).min()
        for dt in arr_types:
            column = getattr(result, dt.__name__)
            is_correct = bool(np.all(column == np.array(correct, dtype=dt)))
            self.assertTrue(
                is_correct,
                msg=f"Incorrect result in column of datatype {dt} when filter was applied on sorted groupby. {column} {correct}",
            )

    def test_math_ops_float_return(self):
        gb = ds.gb("bool")
        result_dict = {
            "mean": [2.5, 3.333333333],
            "std": [2.121320343560, 1.527525231652],
            "median": [2.5, 3],
            "var": [4.5, 2.333333333333],
        }
        for name, correct in result_dict.items():
            call = getattr(gb, name)
            ds_result = call()
            for dt in arr_types:
                column = getattr(ds_result, dt.__name__)
                for idx, item in enumerate(column):
                    is_accurate = isclose(item, correct[idx], rel_tol=1e-06)
                    self.assertTrue(
                        is_accurate,
                        msg=f"Incorrect result for keys {ds_result.label_get_names()} in column with datatype {dt} after {name} operation. {item} != {correct[idx]}",
                    )

            # make sure string columns are not included
            for dt in arr_types_string:
                was_included = hasattr(ds_result, dt.__name__)
                self.assertFalse(
                    was_included,
                    msg=f"String-like column {dt.__name__} was included in {name} math operation.",
                )

    def test_keys_included(self):
        """
        By default, groupby keys will not be included in the dataset's main data for the groupby calculation.
        The keyword include_keys can be set to include them.
        """
        # now included
        ds_result = ds.gb("bool").sum()
        was_included = hasattr(ds_result, "bool")
        self.assertTrue(
            was_included,
            msg=f"Groupby key columns were included in main data. (they should not be by default)",
        )

        # included (by keyword)
        ds_result = ds.gb("bool", include_keys=True).sum()
        was_included = hasattr(ds_result, "bool")
        self.assertTrue(
            was_included,
            msg=f"Groupby key columns were NOT included in main data, even though the keyword was set.",
        )

        # included (by global flag)
        GroupBy.include_keys = True
        ds_result = ds.gb("bool").sum()
        was_included = hasattr(ds_result, "bool")
        self.assertTrue(
            was_included,
            msg=f"Groupby key columns were NOT included in main data, even though the global flag was set.",
        )
        GroupBy.include_keys = False

    def test_ops_str_allowed(self):
        gb = ds.gb("bool")
        result_dict = {"first": [4, 3], "last": [1, 2]}
        for name, correct in result_dict.items():
            call = getattr(gb, name)
            ds_result = call()
            # check numeric columns
            for dt in arr_types:
                column = getattr(ds_result, dt.__name__)
                is_correct = bool(np.all(column == np.array(correct, dtype=dt)))
                self.assertTrue(
                    is_correct,
                    msg=f"Incorrect result for keys {ds_result.label_get_names()} in column with datatype {dt} after {name} operation.",
                )

            # check string columns
            for dt_str in arr_types_string:
                column = getattr(ds_result, dt_str.__name__).astype(np.str_)
                is_correct = bool(np.all(column == np.array(correct, dtype=np.str_)))
                self.assertTrue(
                    is_correct,
                    msg=f"Incorrect result for keys {ds_result.label_get_names()} in column with datatype {dt_str} after {name} operation.",
                )

    def test_math_ops_gb_nums_same_return(self):
        GroupBy.include_keys = True
        test_funcs = ["sum", "first", "last", "nansum"]
        for name in test_funcs:
            # check sorted
            for dt in arr_types:
                call = getattr(ds.gb(dt.__name__), name)
                ds_result = call()
                for dt2 in arr_types:
                    column = getattr(ds_result, dt2.__name__)
                    is_correct = bool(np.all(column == np.array(num_list_sorted, dtype=dt)))
                    self.assertTrue(
                        is_correct,
                        msg=f"Incorrect result for column {dt2.__name__} after grouping by {dt.__name__} for {name} operation. (sorted)",
                    )

            # check unsorted
            for dt in arr_types:
                call = getattr(ds.gbu(dt.__name__), name)
                ds_result = call()
                for dt2 in arr_types:
                    column = getattr(ds_result, dt2.__name__)
                    is_correct = bool(np.all(column == np.array(num_list, dtype=dt)))
                    self.assertTrue(
                        is_correct,
                        msg=f"Incorrect result for column {dt2.__name__} after grouping by {dt.__name__} for {name} operation. (unsorted)",
                    )
        GroupBy.include_keys = False

    # funcs that return nan on single item
    def test_math_ops_nan_return(self):
        GroupBy.include_keys = True
        test_funcs = ["std", "var", "nanstd", "nanvar"]
        for dt in arr_types + arr_types_string:
            gb = ds.gb(dt.__name__)
            # print("***nan_return\n", ds, dt.__name__)
            for func_name in test_funcs:
                call = getattr(gb, func_name)
                ds_result = call()
                for dt2 in arr_types:
                    column = getattr(ds_result, dt2.__name__)
                    # skip over what we grouped by
                    if dt.__name__ != dt2.__name__:
                        for item in column:
                            self.assertNotEqual(
                                item,
                                item,
                                msg=f"Item was not NaN for groupby {dt} in column {dt2} for {func_name} operation. Got {item} instead.",
                            )
        GroupBy.include_keys = False

    def test_count(self):
        """
        The result of count only returns a single column, which displays a count of the number of keys in that grouping.
        """
        result_ds = ds.gb("bool").count()
        self.assertEqual(result_ds._ncols, 2, msg=f"Groupby count did not return a single column.")

        correct_count = np.array([2, 3])
        result_arr = result_ds.Count
        is_correct = bool(np.all(correct_count == result_arr))
        self.assertTrue(
            is_correct,
            msg=f"Groupby bool did not product the correct result for count. returned {result_arr}, expected {correct_count}",
        )

    def test_nan_funcs(self):
        GroupBy.include_keys = True
        nan_dict = {
            "f32": np.arange(5, dtype=np.float32),
            "f64": np.arange(5, dtype=np.float64),
        }
        nan_dict["f32"][0] = np.nan
        nan_dict["f64"][0] = np.nan
        nan_ds = Dataset(nan_dict)
        # TODO: add when more nan functions have been implemented
        correct = np.array([1.0, 2.0, 3.0, 4.0, 0.0])
        nan_funcs = ["nansum"]
        col_names = ["f32", "f64"]

        for c_name in col_names:
            gb = nan_ds.gb(c_name)
            for f_name in nan_funcs:
                call = getattr(gb, f_name)
                result_ds = call()
                for result_name in col_names:
                    column = getattr(result_ds, result_name)
                    # skip groupby column
                    if c_name != result_name:
                        is_correct = bool(np.all(column == correct))
                        self.assertTrue(
                            is_correct,
                            msg=f"Incorrect result for groupby {c_name} in column {result_name} for {f_name} operation.",
                        )
        GroupBy.include_keys = False

    def test_not_supported(self):
        """
        Remove these from the following list when they are implemented.
        """
        not_implemented = [
            "sem",
            "ohlc",
            "describe",
            "expanding",
            "ngroup",
            "rank",
            "head",
            "tail",
            "stack",
            "unstack",
        ]
        # 8/28/2018 - SJK removed nth
        not_implemented_one_arg = ["resample"]
        not_implemented_two_arg = ["transform"]
        gb = ds.gb("bool")

        for func in not_implemented:
            call = getattr(gb, func)
            with self.assertRaises(
                NotImplementedError,
                msg=f"Failed to raise an error for {func}. If this HAS been implemented, remove test from list in test_not_supported in test_sfw_groupby.py",
            ):
                result = call()

        for func in not_implemented_one_arg:
            call = getattr(gb, func)
            with self.assertRaises(
                NotImplementedError,
                msg=f"Failed to raise an error for {func}. If this HAS been implemented, remove test from list in test_not_supported in test_sfw_groupby.py",
            ):
                result = call(1)

        for func in not_implemented_two_arg:
            with self.assertRaises(
                NotImplementedError,
                msg=f"Failed to raise an error for {func}. If this HAS been implemented, remove test from list in test_not_supported in test_sfw_groupby.py",
            ):
                result = call(1, 2)

    def array_equal(self, arr1, arr2):
        subr = arr1 - arr2
        sumr = sum(subr == 0)
        result = sumr == len(arr1)
        if not result:
            print("array comparison failed", arr1, arr2)
        return result

    def assertArrayEqual(self, ar1, ar2):
        assert (ar1 == ar2).all(axis=None), f"array equality failed\nar1 {ar1}\nar2 {ar2}"

    def assertArrayAlmostEqual(self, ar1, ar2, places=5):
        for v1, v2 in zip(ar1, ar2):
            self.assertAlmostEqual(v1, v2, places)

    def test_gb_categoricals(self):
        codes = [1, 44, 44, 133, 75, 75, 75, 1]
        stringlist = ["a", "b", "c", "d", "e", "e", "f", "g"]
        c1 = Categorical(codes, LikertDecision, sort_gb=True)
        c2 = Categorical(stringlist)
        d = {"nums": np.arange(8)}

        # from enum only
        d_enum = d.copy()
        d_enum["cat_from_enum"] = c1
        ds_enum = Dataset(d_enum)
        enum_result = ds_enum.gb("cat_from_enum").sum()
        correct = FastArray([3, 15, 3, 7], dtype=np.int64)
        self.assertTrue(
            self.array_equal(correct, enum_result.nums),
            msg=f"Incorrect sum when grouping by enum categorical.\nExpected {correct}\nActual {enum_result.nums}",
        )

        # from list only
        d_list = d.copy()
        d_list["cat_from_list"] = c2
        ds_list = Dataset(d_list)
        list_result = ds_list.gb("cat_from_list").sum()
        correct = FastArray([0, 1, 2, 3, 9, 6, 7], dtype=np.int64)
        self.assertTrue(
            self.array_equal(correct, list_result.nums),
            msg=f"Incorrect sum when grouping by list categorical.",
        )

        d_both = d_enum.copy()
        d_both["cat_from_list"] = c2
        ds_both = Dataset(d_both)

        # by enum, list
        result = ds_both.gb(["cat_from_enum", "cat_from_list"]).sum()
        num_result = result.nums
        correct = FastArray([0, 7, 1, 2, 9, 6, 3], dtype=np.int64)
        self.assertTrue(
            self.array_equal(correct, num_result),
            msg=f"Incorrect sum when grouping by enum, list categoricals.",
        )

        # by list, enum
        result = ds_both.gb(["cat_from_list", "cat_from_enum"]).sum()
        num_result = result.nums
        correct = FastArray([0, 1, 2, 3, 9, 6, 7], dtype=np.int64)
        self.assertTrue(
            self.array_equal(correct, num_result),
            msg=f"Incorrect sum when grouping by list, enum categoricals.",
        )

    def test_ops(self):
        ds = Dataset(
            {
                "test": arange(300000) % 3,
                "test2": arange(300000.0),
                "test2i": arange(300000),
                "test3": arange(300000) % 3,
            }
        )
        gb = ds.groupby("test")
        result = gb.mean()
        self.assertTrue(result.test2[0] == result.test2i[0])
        self.assertTrue(result.test2[1] == result.test2i[1])
        self.assertTrue(result.test3[1] == 1.0)
        result = gb.median()
        result = gb.trimbr()
        result = gb.nanmedian()
        result = gb.nanquantile(q=0.1)
        result = gb.rolling_quantile(q=0.71)
        result = gb.nanquantile(q=[0.15, 0.45, 0.97])
        result = gb.percentile(q=[0.0, 50.0, 0.6, 100.0])

    def test_projections(self):
        num_rows_trade = 1_000_000
        num_symbols = 450
        Trade_Dates = ["20180602", "20180603", "20180604", "20180605", "20180606"]
        Exchanges = np.array(["EXCH1", "EXCH2", "EXCH3"])
        np.random.seed(1234)
        ds = Dataset(
            {
                "SymbolID": np.random.randint(0, num_symbols, size=num_rows_trade),
                "Exchange": Exchanges[np.random.randint(0, Exchanges.shape[0], size=num_rows_trade)],
                "Trade_Date": [Trade_Dates[int(i * len(Trade_Dates) / num_rows_trade)] for i in range(num_rows_trade)],
                "Time": [int(i % (num_rows_trade / len(Trade_Dates))) for i in range(num_rows_trade)],
                "Price": 100 * (1.0 + 0.0005 * np.random.randn(num_rows_trade)),
                "Size": 10 * np.array(1 + 30 * np.random.rand(num_rows_trade), dtype=np.int64),
            }
        )
        num_rows_quote = 1_000_000
        ds2 = Dataset(
            {
                "SymbolID": np.random.randint(0, num_symbols, size=num_rows_quote),
                "Exchange": Exchanges[np.random.randint(0, Exchanges.shape[0], size=num_rows_quote)],
                "Trade_Date": [Trade_Dates[int(i * len(Trade_Dates) / num_rows_quote)] for i in range(num_rows_quote)],
                "Time": [int(i % (num_rows_quote / len(Trade_Dates))) for i in range(num_rows_quote)],
                "Bid": 100 * (1.0 - 0.001 + 0.0005 * np.random.randn(num_rows_quote)),
                "Ask": 100 * (1.0 + 0.001 + 0.0005 * np.random.randn(num_rows_quote)),
            }
        )
        threshold = Dataset({"Is_Below_Thresdhold": np.random.rand(num_rows_quote) < 0.75})
        trade_time = Dataset({"time_2500": (ds.Time / 2500).astype(int)})
        trades = Dataset({}).concat_columns([ds, trade_time], do_copy=False)

        # Create GroupBy and corresponding Categorical
        trade_gb = trades.groupby(["SymbolID", "Exchange", "Trade_Date", "time_2500"])
        trade_cat = Categorical([ds.SymbolID, ds.Exchange, ds.Trade_Date, trade_time.time_2500])

        # Call sum() and count()
        self.assertEqual(trade_gb.sum().shape, (455654, 7))
        self.assertEqual(trade_cat.sum(ds).shape, (455654, 7))
        self.assertEqual(trade_gb.count().shape, (455654, 5))
        # 8/24/2018 SJK - multikey categorical groupby now returns multiple columns for groupby keys
        self.assertEqual(trade_cat.count().shape, (455654, 5))
        b1 = trade_gb.count().Count.mean()
        b1c = trade_cat.count().Count.mean()
        b2 = trade_gb.count().shape[0]
        self.assertAlmostEqual(ds.shape[0], b1 * b2, places=5)
        self.assertAlmostEqual(ds.shape[0], b1c * b2, places=5)

        # Create ds augmented with filtered ID
        trade_ds = Dataset({"ID": trade_gb.grouping.ikey})
        trade_ds_below_threshold = ds * threshold.Is_Below_Thresdhold
        trade_ds_below_thresholdb = Dataset.concat_columns([trade_ds_below_threshold, trade_ds], do_copy=False)

        # Create trade_ds size projection using GroupBy
        trade_gb_id = trade_ds_below_thresholdb.groupby("ID")
        trade_sizes_ds = trade_gb_id["Size"].sum()
        trade_size_ds = trade_sizes_ds.Size[trade_ds_below_thresholdb.ID - 1]
        self.assertEqual(trade_size_ds.shape[0], ds.shape[0])

        # Create trade_ds size projection using Categorical
        trade_sizes_cat_ds = trade_cat.sum(trade_ds_below_thresholdb.Size)
        trade_size_cat_ds = trade_sizes_cat_ds.Size[trade_cat - 1]
        self.assertArrayAlmostEqual(trade_size_ds, trade_size_cat_ds, places=6)

        # Create trade_ds size projection using Pandas groupby
        ptrade_ds_below_thresholdb = dataset_as_pandas_df(trade_ds_below_thresholdb)
        ptrade_gb_id = ptrade_ds_below_thresholdb.groupby("ID")
        trade_sizes_pd_ds = ptrade_gb_id.sum()
        trade_size_pd_ds = trade_sizes_pd_ds.Size.values[ptrade_gb_id.ngroup()]
        self.assertArrayAlmostEqual(trade_size_ds, trade_size_pd_ds, places=6)

    def test_reductions(self):
        message_types = [
            "CREATE",
            "RUN",
            "CREATE",
            "RUN",
            "RUN",
            "RUN",
            "RUN",
            "CANCEL",
            "RUN",
            "RUN",
            "RUN",
            "CANCEL",
        ]
        order_ids = [1, 1, 2, 1, 2, 2, 1, 2, 2, 2, 2, 1]
        seconds = [50, 70, 72, 75, 90, 88, 95, 97, 98, 115, 116, 120]
        shares = [0, 200, 0, 500, 100, 400, 100, 0, 300, 150, 150, 0]
        d2 = dict(
            message_type=message_types,
            order_id=order_ids,
            second=seconds,
            shares=shares,
        )
        dat = Dataset(d2)
        dat = dat[["order_id", "message_type", "second", "shares"]]

        # Numeric reduction
        dsr = dat.groupby("order_id").sum()
        self.assertEqual(dsr.shape, (2, 3))
        self.assertArrayEqual(dsr.order_id, [1, 2])
        self.assertArrayEqual(dsr.second, [410, 676])
        self.assertArrayEqual(dsr.shares, [800, 1100])

        # Numeric reduction with all columns returned
        dsr = dat.groupby("order_id", return_all=True).sum()
        self.assertEqual(dsr.shape, (2, 4))
        self.assertEqual(dsr.keys()[1], "message_type")

        # Order-based reduction
        dsr = dat.groupby("order_id").first()
        self.assertEqual(dsr.shape, (2, 4))
        self.assertArrayEqual(dsr.order_id, [1, 2])
        self.assertArrayEqual(dsr.message_type, ["CREATE", "CREATE"])
        self.assertArrayEqual(dsr.second, [50, 72])
        self.assertArrayEqual(dsr.shares, [0, 0])

        # Order-based reduction, which returns all columns regardless
        dsr = dat.groupby("order_id", return_all=True).first()
        self.assertEqual(dsr.shape, (2, 4))

        # Order-based reduction with multiple keys
        dsr = dat.groupby(["order_id", "message_type"]).first()
        self.assertEqual(dsr.shape, (6, 4))
        self.assertArrayEqual(dsr.order_id, [1, 1, 1, 2, 2, 2])
        self.assertArrayEqual(dsr.message_type, ["CANCEL", "CREATE", "RUN", "CANCEL", "CREATE", "RUN"])
        self.assertArrayEqual(dsr.second, [120, 50, 70, 97, 72, 90])
        self.assertArrayEqual(dsr.shares, [0, 0, 200, 0, 0, 100])

        # On a subset of columns
        gb = dat.groupby("order_id")
        dsr = gb["shares"].sum()
        self.assertEqual(dsr.shape, (2, 2))
        self.assertArrayEqual(dsr.shares, [800, 1100])

        # Accumulating function
        dsr = dat.groupby("order_id").cumsum()
        self.assertEqual(dsr.shape, (12, 2))
        self.assertArrayEqual(dsr.shares, [0, 200, 0, 700, 100, 500, 800, 500, 800, 950, 1100, 800])

        # return_all has no effect with accumulating functions
        # 8/23/2018 SJK - changed behavior so return all shows the keys
        dsr = dat.groupby("order_id", return_all=True).cumsum()
        self.assertEqual(dsr.shape, (12, 3))

        # Add cum_shares back to a dataset
        dat["cum_shares"] = dat.groupby("order_id").shares.cumsum().shares
        self.assertEqual(dat.shape, (12, 5))
        self.assertArrayEqual(dat.cum_shares, gb.shares.cumsum().shares)

        # On a subset of columns
        dsr = dat.groupby("order_id")[["shares", "second"]].cumsum()
        self.assertEqual(dsr.shape, (12, 2))
        self.assertArrayEqual(dsr.shares, [0, 200, 0, 700, 100, 500, 800, 500, 800, 950, 1100, 800])
        self.assertArrayEqual(dsr.second, [50, 120, 72, 195, 162, 250, 290, 347, 445, 560, 676, 410])

        # On a subset of columns with a filter
        f = FastArray(
            [
                True,
                False,
                True,
                False,
                True,
                False,
                True,
                False,
                True,
                False,
                True,
                False,
            ]
        )
        dsr = dat.groupby("order_id")[["shares", "second"]].cumsum(filter=f)
        self.assertEqual(dsr.shape, (12, 2))
        self.assertArrayEqual(dsr.shares, [0, 0, 0, 0, 100, 100, 100, 100, 400, 400, 550, 100])
        self.assertArrayEqual(dsr.second, [50, 50, 72, 50, 162, 162, 145, 162, 260, 260, 376, 145])

        # On shares and second with filter at groupby construction
        dsr = dat.groupby("order_id", filter=f)[["shares", "second"]].cumsum()
        inv = INVALID_DICT[dsr.shares[0].dtype.num]
        self.assertEqual(dsr.shape, (12, 2))
        self.assertArrayEqual(dsr.shares, [0, inv, 0, inv, 100, inv, 100, inv, 400, inv, 550, inv])
        self.assertArrayEqual(dsr.second, [50, inv, 72, inv, 162, inv, 145, inv, 260, inv, 376, inv])

        # Using agg function
        dsr = gb[["second", "shares"]].agg(["sum", "mean"])
        self.assertEqual(dsr.shape, (2, 2))
        self.assertArrayEqual(dsr.Sum.second, [410, 676])
        self.assertArrayEqual(dsr.Sum.shares, [800, 1100])
        self.assertArrayAlmostEqual(dsr.Mean.second, [82.00, 96.57], places=2)
        self.assertArrayAlmostEqual(dsr.Mean.shares, [160.00, 157.14], places=2)

        # Check for issue when bracket indexing on groupby
        f = open(os.devnull, "w")
        print(gb, file=f)
        f.close()
        dsr = gb[["second", "shares"]].agg(["sum", "mean"])

        # Using different functions on different columns
        dsr = gb.agg({"second": "sum", "shares": ["max", "mean"]})
        self.assertEqual(dsr.shape, (2, 3))
        self.assertArrayEqual(dsr.Sum.second, [410, 676])
        self.assertArrayEqual(dsr.Max.shares, [500, 400])
        self.assertArrayAlmostEqual(dsr.Mean.shares, [160.00, 157.14], places=2)

        # Using numpy functions
        dsr = gb.agg({"second": np.sum, "shares": [np.max, np.mean]})
        self.assertEqual(dsr.shape, (2, 3))
        self.assertArrayEqual(dsr.Sum.second, [410, 676])
        self.assertArrayEqual(dsr.Max.shares, [500, 400])
        self.assertArrayAlmostEqual(dsr.Mean.shares, [160.00, 157.14], places=2)

        # Alternate way to add to multiset
        gb = dat.groupby("order_id")
        ms = gb[["shares"]].agg(["max", "mean"])
        ms.Sum = gb[["second"]].sum()
        self.assertEqual(ms.shape, (2, 3))
        self.assertArrayEqual(ms.Sum.second, [410, 676])
        self.assertArrayEqual(ms.Max.shares, [500, 400])
        self.assertArrayAlmostEqual(ms.Mean.shares, [160.00, 157.14], places=2)

    def test_iter(self):
        correct_keys = FastArray(["e", "d", "b", "c", "a"])
        correct_idx = [[0, 1, 4, 7], [2, 9], [3], [5, 6], [8]]
        str_arr = FastArray(["e", "e", "d", "b", "e", "c", "c", "e", "a", "d"])

        gb = Dataset({"keycol": str_arr, "idxcol": arange(10)})
        gb = gb.gb("keycol")
        for i, tup in enumerate(gb):
            self.assertEqual(tup[0], correct_keys[i])
            self.assertTrue(bool(np.all(tup[1].idxcol == correct_idx[i])))

    def test_shift(self):
        import pandas as pd

        ds = Dataset({"col_" + str(i): np.random.rand(30) for i in range(5)})
        ds.keycol = np.random.choice(["a", "b", "c"], 30)
        df = pd.DataFrame(ds.asdict())

        rt_result = ds.gb("keycol").shift(periods=-2).trim()
        pd_result = df.groupby("keycol").shift(periods=-2).dropna(axis="rows")

        for k, v in rt_result.items():
            self.assertTrue(bool(np.all(v == pd_result[k])))

        rt_result = ds.gb("keycol").shift(periods=3).trim()
        pd_result = df.groupby("keycol").shift(periods=3).dropna(axis="rows")

        for k, v in rt_result.items():
            self.assertTrue(bool(np.all(v == pd_result[k])))

    def test_diff(self):
        import pandas as pd

        ds = Dataset({"col_" + str(i): np.random.rand(10) for i in range(5)})
        ds.keycol = np.random.choice(["a", "b", "c"], 10)
        df = pd.DataFrame(ds.asdict())

        rt_result = ds.gb("keycol").rolling_diff()
        pd_result = df.groupby("keycol").diff()

        for k, v in rt_result.items():
            pdc = pd_result[k]
            pdcnan = isnan(pdc)
            self.assertTrue(bool(np.all(isnan(v) == pdcnan)), msg=f"{v} {pdc}")

            masked_valid_pd = isnotnan(pdc)
            masked_valid_rt = isnotnan(v)

            self.assertTrue(bool(np.all(masked_valid_pd == masked_valid_rt)))

    def test_pad(self):
        arrsize = 100
        numrows = 20
        ds = Dataset({"time": arange(arrsize * 1.0)})
        ds.data = np.random.randint(numrows, size=arrsize)
        ds.data2 = np.random.randint(numrows, size=arrsize)
        symbols = ["ZYGO", "YHOO", "FB", "GOOG", "IBM"]
        ds.symbol2 = Cat(1 + ds.data, list("ABCDEFGHIJKLMNOPQRST"))
        ds.symbol = Cat(1 + arange(arrsize) % len(symbols), symbols)
        ds.time[[3, 4, 7]] = nan
        newds = ds.gb("symbol").pad()
        self.assertTrue(newds.time[7] == 2.00)
        self.assertTrue(newds.time[3] != newds.time[3])
        newds = ds.gb("symbol").backfill()
        self.assertTrue(newds.time[7] == 12.00)

        # see if we can pull a group
        newds = ds.gb("symbol").get_group("YHOO")
        self.assertTrue(np.all(newds.symbol == "YHOO"))

    def test_transform(self):

        arrsize = 200
        numrows = 7

        ds = Dataset({"time": arange(arrsize * 1.0)})
        ds.data = np.random.randint(numrows, size=arrsize)
        ds.data2 = np.random.randint(numrows, size=arrsize)
        symbols = ["AAPL", "AMZN", "FB", "GOOG", "IBM"]
        ds.symbol = Cat(1 + arange(arrsize) % len(symbols), symbols)
        newds = ds.gb("symbol")["data"].sum(transform=True)

        # removed from test since gbkeys not returned in transform
        # self.assertTrue(np.all(newds.symbol == ds.symbol))
        catds = ds.symbol.sum(ds.data, transform=True)
        self.assertTrue(np.all(newds[0] == catds[0]))
        # test showfilter
        catds = ds.symbol.sum(ds.data, showfilter=True, transform=True)
        self.assertTrue(np.all(newds[0] == catds[0]))

        # test diff
        result1 = ds.gb("symbol").apply_nonreduce(TypeRegister.FastArray.diff)
        result2 = ds.gb("symbol").diff()
        self.assertTrue(result1.equals(result2))


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