# $Id: //Depot/Source/SFW/riptable/Python/core/riptable/tests/test_multiset.py#4 $
import unittest

import numpy as np

# from io import StringIO
from riptable import Dataset, Multiset, Struct


class Multiset_Test(unittest.TestCase):
    @staticmethod
    def get_basic_dataset(keyfield_value=None, nr=10):
        ds = Dataset({_k: list(range(_i * nr, (_i + 1) * nr)) for _i, _k in enumerate("abcdefghijklmnop")})
        if keyfield_value is not None:
            ds.keyfield = keyfield_value
        return ds

    def test_ctor_01(self):
        dd = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("bac")}
        ms = Multiset(dd)
        self.assertEqual(list(ms.keys()), ["b", "a", "c"])
        for _i, _k in enumerate("bac"):
            self.assertEqual(ms[_k].keyfield[0], _i)
            self.assertEqual(getattr(ms, _k).keyfield[0], _i)
        with self.assertRaises(TypeError):
            _ = Multiset(ms)
        with self.assertRaises(TypeError):
            _ = Multiset(list(dd.items()))
        assert Multiset(None).shape == (
            0,
            0,
        ), f"Multiset default argument for 'None' input value should create an empty Multiset."
        ms = Multiset()
        self.assertEqual(ms.shape, (0, 0))
        ms = Multiset({})
        self.assertEqual(ms.shape, (0, 0))

    def test_ctor_02(self):
        dd = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("bac")}
        dd["d"] = self.get_basic_dataset(3)
        ms = Multiset(dd)
        self.assertEqual(list(ms.keys()), ["b", "a", "c", "d"])
        dd["d"] = Multiset({_k: self.get_basic_dataset(_i) for _i, _k in enumerate("bac")})
        ms = Multiset(dd)
        self.assertEqual(list(ms.keys()), ["b", "a", "c", "d"])
        # Number of rows in contained datasets must match
        dd["d"] = self.get_basic_dataset(3, nr=9)
        with self.assertRaises(ValueError):
            _ = Multiset(dd)
        # Number of rows in contained datasets and multisets must match
        dd["d"] = Multiset({})
        with self.assertRaises(ValueError):
            _ = Multiset(dd)
        dd["d"] = Multiset({_k: self.get_basic_dataset(_i, nr=9) for _i, _k in enumerate("bac")})
        with self.assertRaises(ValueError):
            _ = Multiset(dd)
        # Type of input must be valid
        dd["d"] = Struct({})
        with self.assertRaises(TypeError):
            _ = Multiset(dd)
        dd["d"] = None
        with self.assertRaises(TypeError):
            _ = Multiset(dd)
        dd["d"] = 1
        with self.assertRaises(TypeError):
            _ = Multiset(dd)

    def test_modify(self):
        dd = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("bac")}
        ms = Multiset(dd)
        self.assertEqual(list(ms.keys()), list("bac"))
        with self.assertRaises(ValueError):
            ms.d = Dataset({})
        ms.d = Dataset({"a": range(10)})
        self.assertEqual(list(ms.keys()), list("bacd"))
        with self.assertRaises(TypeError):
            ms.e = Struct({})
        with self.assertRaises(TypeError):
            ms.d = Struct({})
        self.assertEqual(list(ms.keys()), list("bacd"))
        self.assertIsInstance(ms.d, Dataset)
        ms._lock()
        with self.assertRaises(AttributeError):
            ms.e = Dataset({})
        ms._unlock()
        with self.assertRaises(ValueError):
            ms.e = Dataset({})
        ms.e = Dataset({"a": range(10)})
        self.assertEqual(list(ms.keys()), list("bacde"))
        self.assertEqual(ms.e.shape, (10, 1))
        self.assertIs(ms.e, ms[4])
        self.assertEqual(ms[4].shape, (10, 1))
        # Added Dataset must have same number of rows
        with self.assertRaises(ValueError):
            ms["f"] = Dataset({"a": 1, "b": 2})
        ms["f"] = Dataset({"a": range(10), "b": range(10)})
        ms[5] = Dataset({"a": [3] * 10, "b": [4] * 10})
        with self.assertRaises(IndexError):
            ms[6] = Dataset({"a": [3] * 10, "b": [4] * 10})
        self.assertEqual(list(ms.keys()), list("bacdef"))
        with self.assertRaises(IndexError):
            ms[["e", "g"]] = [Dataset({"a": 5, "b": 6}), Dataset({"a": 7, "b": 8})]
        ms[["e", "f"]] = [
            Dataset({"a": [5] * 10, "b": [6] * 10}),
            Dataset({"a": [7] * 10, "b": [8] * 10}),
        ]

    def test_multiget(self):
        # First testing case where all the contained Datasets are different:
        dd1 = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("ABC")}
        dd1["A"].col_remove(list("abc"))
        dd1["B"].col_remove(list("def"))
        dd1["D"] = Multiset({"SUB": dd1["C"]})
        ms1 = Multiset(dd1)
        ms2 = ms1.multiget(list("aeg"))
        self.assertEqual(list(ms1.keys()), list("ABCD"))
        self.assertEqual(list(ms2.keys()), list("ABC"))
        self.assertEqual(list(ms2.A.keys()), list("eg"))
        self.assertEqual(list(ms2.B.keys()), list("ag"))
        self.assertEqual(list(ms2.C.keys()), list("aeg"))
        ms3 = ms1.multiget(slice(2, 15, 4))
        self.assertEqual(list(ms3.keys()), list("ABC"))
        self.assertEqual(list(ms3.A.keys()), list("fjn"))
        self.assertEqual(list(ms3.B.keys()), list("cjn"))
        self.assertEqual(list(ms3.C.keys()), list("cgko"))
        # Now, when they all match:
        dd2 = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("ABCD")}
        ms4 = Multiset(dd2)
        self.assertEqual(list(ms4.keys()), list("ABCD"))
        ms5 = ms4.multiget(list("cgko"))
        ms6 = ms4.multiget(slice(2, 15, 4))
        filt = np.zeros((ms4.A.get_ncols(),), dtype=bool)
        filt[2:15:4] = True
        ms7 = ms4.multiget(filt)
        for ms in (ms5, ms6, ms7):
            self.assertEqual(list(ms.keys()), list("ABCD"))
            for dskey in "ABCD":
                self.assertEqual(list(getattr(ms, dskey).keys()), list("cgko"))
        ms8 = ms4.multiget("c")
        ms9 = ms4.multiget(["c"])
        for ms in (ms8, ms9):
            self.assertEqual(list(ms.keys()), list("ABCD"))
            for dskey in "ABCD":
                self.assertEqual(list(getattr(ms, dskey).keys()), list("c"))

    def test_copy(self):
        dd = {_k: self.get_basic_dataset(_i) for _i, _k in enumerate("bac")}
        ms = Multiset(dd)
        ms1 = ms.copy()
        self.assertEqual(list(ms1.keys()), ["b", "a", "c"])
        self.assertTrue((ms1.a == ms.a).all(axis=None))
        self.assertTrue((ms1.b == ms.b).all(axis=None))
        self.assertTrue((ms1.c == ms.c).all(axis=None))
        self.assertFalse(ms1.is_locked())
        ms._lock()
        ms2 = ms.copy()
        self.assertTrue(ms2.is_locked())

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

    def test_flatten(self):
        accuracy = 7
        message_types = [
            "NEW",
            "EXECUTE",
            "NEW",
            "EXECUTE",
            "EXECUTE",
            "EXECUTE",
            "EXECUTE",
            "CANCEL",
            "EXECUTE",
            "EXECUTE",
            "EXECUTE",
            "CANCEL",
        ]
        order_ids = [1, 1, 2, 1, 2, 2, 1, 2, 2, 2, 2, 1]
        milliseconds = [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]
        dat = Dataset(
            dict(
                message_type=message_types,
                order_id=order_ids,
                millisecond=milliseconds,
                shares=shares,
            )
        )
        dat = dat[["order_id", "message_type", "millisecond", "shares"]]
        gb = dat.groupby("order_id")
        ms1 = gb[["millisecond", "shares"]].aggregate(["sum", "mean"])
        # Flatten horizontally
        f1 = ms1.flatten()
        self.assertEqual(
            f1.keys(),
            [
                "order_id",
                "Sum_millisecond",
                "Sum_shares",
                "Mean_millisecond",
                "Mean_shares",
            ],
        )
        self.assertTrue((f1.Sum_millisecond == [410, 676]).all(axis=None))
        # Flatten vertically
        f2 = ms1.flatten(horizontal=False)
        self.assertEqual(f2.keys(), ["order_id", "Column", "millisecond", "shares"])
        self.almost_eq(f2.millisecond, [410.0, 676.0, 82.0, 96.57142857], accuracy)
        ms2 = gb[["millisecond", "shares"]].aggregate(["min", "max"])
        # Flatten multiset containing multisets
        ms3 = Multiset({"ms1": ms1, "ms2": ms2})
        f3 = ms3.flatten()
        self.assertEqual(
            f3.keys(),
            [
                "order_id",
                "ms1_Sum_millisecond",
                "ms1_Sum_shares",
                "ms1_Mean_millisecond",
                "ms1_Mean_shares",
                "ms2_Min_millisecond",
                "ms2_Min_shares",
                "ms2_Max_millisecond",
                "ms2_Max_shares",
            ],
        )
        self.assertTrue((f3.ms1_Sum_millisecond == [410, 676]).all(axis=None))
        f3v = ms3.flatten(horizontal=False)
        self.assertTrue(
            (
                f3v.Column
                == [
                    b"ms1_Sum",
                    b"ms1_Sum",
                    b"ms1_Mean",
                    b"ms1_Mean",
                    b"ms2_Min",
                    b"ms2_Min",
                    b"ms2_Max",
                    b"ms2_Max",
                ]
            ).all(axis=None)
        )
        # Flatten multiset containing multisets and a dataset
        ds = gb[["millisecond", "shares"]].std()
        ms4 = Multiset({"ms1": ms1, "ms2": ms2, "Std": ds})
        f4 = ms4.flatten()
        f4.label_remove()
        self.assertEqual(
            f4.keys(),
            [
                "order_id",
                "ms1_Sum_millisecond",
                "ms1_Sum_shares",
                "ms1_Mean_millisecond",
                "ms1_Mean_shares",
                "ms2_Min_millisecond",
                "ms2_Min_shares",
                "ms2_Max_millisecond",
                "ms2_Max_shares",
                "Std_millisecond",
                "Std_shares",
            ],
        )
        self.assertTrue((f4.ms1_Sum_millisecond == [410, 676]).all(axis=None))
        f4 = ms4.flatten(horizontal=False)
        self.assertAlmostEqual(f4[9, "millisecond"], 15.4903964, places=5)
        self.assertAlmostEqual(f4[9, "shares"], 148.4042099, places=5)
        # Flatten multiset containing multisets and a dataset with non-matching column
        ds1 = gb[["shares"]].std()
        ms5 = Multiset({"ms1": ms1, "ms2": ms2, "Std": ds1})
        f5 = ms5.flatten()
        f5.label_remove()
        self.assertEqual(
            f5.keys(),
            [
                "order_id",
                "ms1_Sum_millisecond",
                "ms1_Sum_shares",
                "ms1_Mean_millisecond",
                "ms1_Mean_shares",
                "ms2_Min_millisecond",
                "ms2_Min_shares",
                "ms2_Max_millisecond",
                "ms2_Max_shares",
                "Std_shares",
            ],
        )
        self.assertTrue((f5.ms1_Sum_millisecond == [410, 676]).all(axis=None))
        with self.assertRaises(ValueError):
            ms5.flatten(horizontal=False)


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